I’m new to pytorch lightning and currently training a custom model to classify diseases which had 2 pairs of cnn, maxpool and relu layers. The same model was working very well when using normal pytorch module. I wanted to add more functionality so choose pytorch lightning. The same architecture is now giving error, can someone please guide me through
Normal pytorch module:
class NeuralNetwork(nn.Module):
def init(self):
super(NeuralNetwork, self).init()
self.cnn = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=2)
self.relu = nn.LeakyReLU()
self.maxpool = nn.MaxPool2d(kernel_size=2)
self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=2)
self.relu2 = nn.LeakyReLU()
self.maxpool2 = nn.MaxPool2d(kernel_size=2)
self.fc1 = nn.Linear(326565, 1000)
self.fc2 = nn.Linear(1000, 12)
def forward(self, x):
out = self.cnn(x)
out = self.relu(out)
out = self.maxpool(out)
out = self.cnn2(out)
out = self.relu2(out)
out = self.maxpool2(out)
out = out.view(out.size(0), -1)
out = self.fc1(out)
out = self.fc2(out)
return out
Implementing pytorch lightning:
class NeuralNetwork(pl.LightningModule):
def init(self):
super(NeuralNetwork, self).init()
self.cnn = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=2)
self.relu = nn.LeakyReLU()
self.maxpool = nn.MaxPool2d(kernel_size=2)
self.cnn2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=2)
self.relu2 = nn.LeakyReLU()
self.maxpool2 = nn.MaxPool2d(kernel_size=2)
self.fc1 = nn.Linear(326565, 1000)
self.fc2 = nn.Linear(1000, 14)
self.loss = nn.CrossEntropyLoss()
def forward(self, x):
out = self.cnn(x)
out = self.relu(out)
out = self.maxpool(out)
out = self.cnn2(out)
out = self.relu2(out)
out = self.maxpool2(out)
out = out.view(out.size(0), -1)
out = self.fc1(out)
out = self.fc2(out)
return out