# Why does softmax cause my nn to not converge

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

• you're probably passing probabilities into an implementation of cross-entropy which expects logits – shimao Sep 1 '20 at 19:25

## 1 Answer

The reason that this didn't work is Pytorch's implementation of cross entropy loss in nn.CrossEntropyLoss expects logits, not the probabilities output by softmax as suggested in shimao's comment.