Hello, I use the following code to define my simple model:
def BlockConv(channels_in, channels_out):
return nn.Sequential(
nn.Conv2d(channels_in, channels_out, (3, 3), padding=1),
nn.MaxPool2d(2),
nn.ReLU()
)
def BlockDeconv(channels_in, channels_out):
return nn.Sequential(
nn.Upsample(scale_factor=2),
nn.Conv2d(channels_in, channels_out, (3, 3), padding=1),
nn.ReLU()
)
def BlockConvWithoutPool(channels_in, channels_out):
return nn.Sequential(
nn.Conv2d(channels_in, channels_out, (3, 3), padding=1)
)
class Colorizer(nn.Module):
def __init__(self):
super().__init__()
self.preconcat = nn.Sequential(
BlockConv(1, 70),
BlockDeconv(70, 64)
)
self.postconcat = nn.Sequential(
BlockConvWithoutPool(65, 3),
nn.Sigmoid()
)
def forward(self, x):
h = self.preconcat(x)
h = torch.cat((h, x), 1)
h = self.postconcat(h)
return h
But it eats all my gpu (80%/8gb) after first iteration of training (128x128 images) on loss.backward()
.
What have I done wrong?
You can check it on Google Colab