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

Question

我在 pytorch 中使用 multiprocessing 包将训练拆分到多个进程。 我的 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)

环境细节：Python-3 和 Linux

Answer 1

他们是一样的。 torch.multiprocessing.Queue内部使用tensor.share_memory_() 。