2017-07-24 74 views
2

我有cntk中的模型,它由一個嵌入層和一個跟隨的LSTM組成。我想將嵌入查找操作放在CPU上,並將其餘的網絡放在GPU上。
在tensorflow,我會做如何在CNTK中爲特定操作設置設備?

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

什麼是cntk等效,我可以包裝成這樣的:

C.layers.Embedding(embedding_size) 

回答

1

在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 
相關問題