I have a pytorch model made up of a several convolutional and groupnorm layers which eventually feed into fully connected and eventually a softmax. With the softmax, the model never converges and plateaus quickly (tried up to epoch > 300). But removing the softmax causes the model to decrease in loss quickly and converge within 50 epochs.
My understanding is that multi-class networks should receive a softmax layer followed by categorical cross entropy loss, why in this case does the model only converge after a fully connected layer?
The inputs to the model are encoded videos of size (128,64,64) all normalized between 0 and 1.
class ActionClass(nn.Module):
def __init__(self, class_num):
super(ActionClass, self).__init__()
self.classNum = class_num
self.conv_b1_1 = nn.Conv2d(161, 128, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b1_1 = nn.GroupNorm(128, 128)
self.conv_b1_2 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b1_2 = nn.GroupNorm(256, 256)
self.conv_b2_1 = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b2_1 = nn.GroupNorm(256, 256)
self.conv_b2_2 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b2_2 = nn.GroupNorm(512, 512)
self.conv_b3_1 = nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b3_1 = nn.GroupNorm(512, 512)
self.conv_b3_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
self.groupNorm_b3_2 = nn.GroupNorm(512, 512)
self.relu = nn.LeakyReLU(inplace=True)
self.softmax = nn.Softmax(dim=0)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, self.classNum)
def forward(self, x):
x = self.conv_b1_1(x)
x = self.groupNorm_b1_1(x)
x = self.relu(x)
x = self.conv_b1_2(x)
x = self.groupNorm_b1_2(x)
x = self.relu(x)
x = self.conv_b2_1(x)
x = self.groupNorm_b2_1(x)
x = self.relu(x)
x = self.conv_b2_2(x)
x = self.groupNorm_b2_2(x)
x = self.relu(x)
x = self.conv_b3_1(x)
x = self.groupNorm_b3_1(x)
x = self.relu(x)
x = self.conv_b3_2(x)
x = self.groupNorm_b3_2(x)
x = self.relu(x)
x = self.avgpool(x)
x = x.flatten(start_dim=1)
x = self.fc(x)
x = self.softmax(x)
return x