I just got a new machine with 2 gtx1080ti, so I wanted to try using nn.DataParallel for faster training. I have created a test code to make sure nn.DataParallel works, but it seems to get stuck in the forward().
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
class RandomDataset(Dataset):
def __init__(self):
self.a = torch.randn(123, 15)
pass
def __getitem__(self, index):
return self.a[index]
def __len__(self):
return 123
class Model(nn.Module):
def __init__(self, i_s, o_s):
super(Model, self).__init__()
self.fc = nn.Linear(i_s, o_s)
def forward(self, x):
output = self.fc(x)
return output
dl = DataLoader(dataset=RandomDataset(), batch_size=16, shuffle=True)
model = Model(15, 3)
print(torch.cuda.device_count())
model = nn.DataParallel(model)
model.to(device)
print('it works!!')
for data in dl:
x = data.to(device)
y = model(x)
print('!!it works')
Training works perfectly if I don’t use nn.DataParallel.
I did not know anything about setting up cuda because I was using cloud service previously, so I just followed the steps here.