Hi,
I have a very deep model and I am trying to implement model sharding as its not working on a single GPU. Code is below:
class ConvolutionalAutoEncoder(nn.Module):
def __init__(self):
super(ConvolutionalAutoEncoder, self).__init__()
self.encoder_block1 = nn.Sequential(
nn.Conv2d(1, 64, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(64, 64, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(64)
)
self.encoder_block2 = nn.Sequential(
nn.Conv2d(64, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.encoder_block3 = nn.Sequential(
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.encoder_block4 = nn.Sequential(
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.Conv2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.decoder_block4 = nn.Sequential(
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.decoder_block3 = nn.Sequential(
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.decoder_block2 = nn.Sequential(
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(128)
)
self.decoder_block1 = nn.Sequential(
nn.ConvTranspose2d(128, 64, 3, stride=1, padding=1),
nn.Tanh(),
nn.ConvTranspose2d(64, 64, 3, stride=1, padding=1),
nn.Tanh(),
nn.BatchNorm2d(64)
)
self.decoder_block0 = nn.Sequential(
nn.ConvTranspose2d(64, 1, 3, stride=1, padding=1),
)
self.encoder_block1 = self.encoder_block1.to('cuda:0')
self.encoder_block2 = self.encoder_block2.to('cuda:0')
self.encoder_block3 = self.encoder_block3.to('cuda:1')
self.encoder_block4 = self.encoder_block4.to('cuda:1')
self.decoder_block4 = self.decoder_block4.to('cuda:2')
self.decoder_block3 = self.decoder_block3.to('cuda:2')
self.decoder_block2 = self.decoder_block2.to('cuda:3')
self.decoder_block1 = self.decoder_block1.to('cuda:3')
self.decoder_block0 = self.decoder_block0.to('cuda:3')
def forward(self, x):
x1 = self.encoder_block1(x)
x2 = self.encoder_block2(x1)
x2 = x2.to('cuda:1')
x3 = self.encoder_block3(x2)
x4 = self.encoder_block4(x3)
x4 = x4.to('cuda:2')
x3 = x3.to('cuda:2')
y4 = self.decoder_block4(x4)
y3 = self.decoder_block3(y4+x3)
y3 = y3.to('cuda:3')
x2 = x2.to('cuda:3')
y2 = self.decoder_block2(y3+x2)
y1 = self.decoder_block1(y2)
y0 = self.decoder_block0(y1)
return y0
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
model = ConvolutionalAutoEncoder()
learning_rate = 0.001
weight_decay = 0.1
momentum = 0.9
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
params = list(model.parameters())
print(len(params))
print(params[0].size()) # conv1's .weight
num_epochs = 30
for epoch in range(num_epochs):
for i, data in enumerate(train_loader):
inp, targ = data
inp = inp.to(device)
targ = targ.to(device)
output = model(inp)
output = output.to(device)
optimizer.zero_grad()
loss = F.binary_cross_entropy_with_logits(output, targ)
loss.backward()
optimizer.step()
if i % 10 == 0:
for param_group in optimizer.param_groups:
print("Current learning rate is: {}".format(param_group['lr']))
print("Epoch[{}/{}]({}/{}): Loss: {:.4f}".format(epoch+1,num_epochs, i, len(train_loader), loss.item()))
The code runs and I think the forward pass works well. I can see the load being passed to the different GPU’s but after one pass it stops working and the code just hangs. I cant even kill it and the only way to get out is to force reboot.
Screenshot of nvidia-smi also attached:
No matter how much I try to kill the process (7711), it does not stop.
Is there something wrong with my code?
Please help.
EDIT:
Ok. So I think I solved the problem. I was computing the loss function in cuda 0 but the last layer was in cuda 3, hence I think it was causing problem while backward propagating (not sure why it didnt throw an error but just froze completely).
So I moved the loss function also to cuda 3 and it worked!!! 
But I dont know why but I am not able to print the loss value i.e.
if i % 10 == 0:
print("Epoch[{}/{}]({}/{}): Loss: {:.4f}".format(epoch+1,num_epochs, i, len(train_loader), loss.item()))
The code’s freezing up in this part. Why is just printing the loss causing a problem?