使用 tensor.share_memory_() 与多处理.跨多个进程训练模型时在 PyTorch 中排队

Question

我正在使用pytorch中的多处理包将训练拆分到多个进程。我的 x 和 y，训练和测试数据是 CUDA 张量。我试图了解使用 tensor.share_memory_(( 和多处理之间的区别。共享 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_()。