Can anyone please explain me how to calculate memory required for this model, and also cuda is saying “Out of memory”, so can there be any optimization or any other suggestions are welcome because i am totally new to ML.
class MODEL(nn.Module) :
def __init__(self) :
super(MODEL, self).__init__()
self.conv1 = nn.Conv2d(1, 64, 3, padding=1)
self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
self.conv4 = nn.Conv2d(256, 512, 3, padding=1)
self.conv5 = nn.Conv2d(512, 256, 3, padding=1)
self.conv6 = nn.Conv2d(256, 128, 3, padding=1)
self.conv7 = nn.Conv2d(128, 64, 3, padding=1)
self.conv8 = nn.Conv2d(64, 32, 3, padding=1)
self.conv9 = nn.Conv2d(32, 2, 3, padding=1)
self.activ = FF.relu
self.pool = nn.AvgPool2d(3, stride=1)
self.padding = nn.ZeroPad2d(1)
self.activ2 = nn.Tanh()
def forward(self, input):
input = self.activ(self.conv1(input))
input = self.activ(self.conv2(input))
input = self.activ(self.conv3(input))
input = self.pool(self.padding(input))
input = self.activ(self.conv4(input))
input = self.activ(self.conv5(input))
input = self.activ(self.conv6(input))
input = self.activ(self.conv7(input))
input = self.pool(self.padding(input))
input = self.activ(self.conv8(input))
input = self.activ(self.conv9(input))
input =self.activ2(input)
input = torch.mul(input, 128.0)
return input
Thank you