gpt4 book ai didi

python - CNTK中如何设置设备进行特定操作?

转载 作者:太空宇宙 更新时间:2023-11-04 05:02:51 24 4
gpt4 key购买 nike

我在 cntk 中有一个模型,它由一个嵌入层和一个后续的 LSTM 组成。我想将嵌入查找操作放在 CPU 上,将网络的其余部分放在 GPU 上。
在 tensorflow 中,我会做

with tf.device("/cpu:0"):

cntk 中我可以将其包装成的等价物是什么:

C.layers.Embedding(embedding_size)

最佳答案

在 CNTK 中,设备控制是针对特定功能/网络的前向/后向传递,而不是单个操作。如果您想在 CPU 中执行特定操作,而在 GPU 中执行其余操作,则必须执行以下操作:

features = C.input_variable(input_dim)
labels = C.input_variable(label_dim)
embedding_input = C.input_variable(embedding_size, needs_gradient=True)
embedding = C.layers.Embedding(embedding_size)(features)
loss = rest_of_your_network(embedding_input, labels)
emb_state, embedding_value = embedding.forward({features: some_data}, keep_for_backward=set(embedding.output), device=C.cpu(), as_numpy=False)
loss_state, loss_value = loss.forward({embedding_input: embedding_value}, keep_for_backward=set(loss.output), device=C.gpu(0), as_numpy=False)
loss_grad_dict = loss.backward(loss_state, {loss.output: np.ones_like(loss_value)}, set(loss.parameters + [embedding_input]))
emb_grad_dict = embedding.backward(emb_state, {embedding.output: loss_grad_dict[embedding_input]}, set(embedding.parameters))
# use these dictionaries to update the parameters with a learner

关于python - CNTK中如何设置设备进行特定操作?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/45291194/

24 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com