Question

我需要转换现有模型以使用DataParallel来在多个GPU上运行，我已经阅读了很长时间的文档，无法了解我需要确切更改的内容，所以我需要一些帮助，谢谢！

class LSTM(nn.Module):
def __init__(self, input_size=1, hidden_layer_size=100, output_size=1):
    super().__init__()
    self.hidden_layer_size = hidden_layer_size

    self.lstm = nn.LSTM(input_size, hidden_layer_size)

    self.linear = nn.Linear(hidden_layer_size, output_size)

    self.hidden_cell = (torch.zeros(1,1,self.hidden_layer_size),
                        torch.zeros(1,1,self.hidden_layer_size))

def forward(self, input_seq):
    lstm_out, self.hidden_cell = self.lstm(input_seq.view(len(input_seq) ,1, -1), self.hidden_cell)
    predictions = self.linear(lstm_out.view(len(input_seq), -1))
    return predictions[-1]


model = LSTM()
loss_function = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)

epochs = 1


 for i in range(epochs):
  count=0
  for seq, labels in train_inout_seq:
    optimizer.zero_grad()
    model.hidden_cell = (torch.zeros(1, 1, model.hidden_layer_size),
                    torch.zeros(1, 1, model.hidden_layer_size))
    y_pred = model(seq)

    single_loss = loss_function(y_pred, labels)
    single_loss.backward()
    optimizer.step()

print("Epoch: %d, loss: %1.5f" % (i, single_loss.item()))

model.eval()

 testPredictions=[]
 realValues=[]
 for seq, labels in test_input_seq:

  with torch.no_grad():
    model.hidden = (torch.zeros(1, 1, model.hidden_layer_size),
                    torch.zeros(1, 1, model.hidden_layer_size))
    testPredictions.append(model(seq).item())
    print(model(seq).item())
    realValues.append(labels)

Pytorch将LSTM模型转换为Dataparallel以用于多种GPU使用

0 个答案: