跨多個進程訓練模型時，在 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_() 。