在多个进程中训练模型时，在PyTorch中使用tensor.share_memory_（）与multiprocessing.Queue

Question

我正在使用pytorch中的多处理程序包将培训划分为多个过程。我的x和y，训练和测试数据是CUDA张量。我试图了解使用tensor.share_memory_（）和multiprocessing.Queue方法共享cuda张量之间的区别。首选哪个，为什么？

这是我当前使用tensor.share_memory_（）的代码。我应该做些什么改变？

def train(model, features, target, epochs=1000):

    X_train, x_test, Y_train, y_test = train_test_split(features,
                                                    target,
                                                    test_size=0.4,
                                                    random_state=0)

    Xtrain_ = torch.from_numpy(X_train.values).float().share_memory_()
    Xtest_ = torch.from_numpy(x_test.values).float().share_memory_()

    Ytrain_ = (torch.from_numpy(Y_train.values).view(1,-1)[0]).share_memory_()
    Ytest_ = (torch.from_numpy(y_test.values).view(1,-1)[0]).share_memory_()


    optimizer = optim.Adam(model.parameters(), lr = 0.01)
    loss_fn = nn.NLLLoss()


    for epoch in range(epochs):

        #training code here

目标方法到此结束

mp.set_start_method('spawn')

model = Net()
model.share_memory()

processes = []

for rank in range(1):
    p = mp.Process(target=train, args=(model, features, target))
    p.start()
    processes.append(p)

Env详细信息：Python-3和Linux

Answer 1

它们是相同的。 torch.multiprocessing.Queue在内部使用tensor.share_memory_()。