Gradient tensor exists for binary output, but is None for multiclass output

ellax123: <p>I am encountering some very strange behavior with my model. I am trying to take the feature extraction layers of ResNet34 and add a final classifier layer onto the end. When I change the final nn.Linear layer from a binary output to a multiclass output, the gradient disappears and my model fails to update its weights.</p>
<p>PyTorch version 1.0.1<br>
Let me know if any additional information is useful, I’ve never made a post like this before.</p>
<p>Example code:</p>
<h3>General setup</h3>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision.models as models

class ResNet34_Model(nn.Module):
    def __init__(self, original_model, num_classes):
        super(TestModel, self).__init__()
            linear_size = 512
            self.features = nn.Sequential(*list(original_model.children())[:-1])
            self.classifier = nn.Sequential(
                nn.Linear(linear_size, num_classes)
            )
        # Freeze those weights
        for p in self.features.parameters():
            p.requires_grad = False

    def forward(self, x):
        f = self.features(x)
        if self.modelName == 'alexnet' :
            f = f.view(f.size(0), 256 * 6 * 6)
        elif self.modelName == 'vgg16':
            f = f.view(f.size(0), -1)
        elif self.modelName == 'resnet' :
            f = f.view(f.size(0), -1)
        elif self.modelName == "densenet":
            # f = f.relu(f, inplace=True)
            # f = f.avg_pool2d(f, kernel_size=7).view(f.size(0), -1)
            f = f.view(f.size(0), -1)
        y = self.classifier(f)
        return y
</code></pre>
<h3>Binary code (gradient is fine)</h3>
<pre><code class="lang-auto">original_model = models.__dict__["resnet34"](pretrained=True)
model = ResNet34_Model(original_model, 1)
model = model.cuda()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = 3e-4, weight_decay=0)
criterion = nn.BCEWithLogitsLoss().cuda()

X_batch = X_batch.cuda() # X_batch and y_batch are pulled from the first iteration of a data loader
y_batch = y_batch.cuda()

# train for 1 epoch
model.train()
output = model(X_batch).squeeze()
loss = criterion(output.float(), y_batch.float())
output_var = torch.sigmoid(output).data
output_class = [ 0 if a&lt;=0.5 else 1  for a in output_var]
optimizer.zero_grad()
loss.backward()
optimizer.step()

# check that gradients exist
for n, p in model.named_parameters():
    print(n)
    if(p.requires_grad) and ("bias" not in n):
        print(p.grad)

</code></pre>
<h4>Output</h4>
<pre><code class="lang-auto">features.0.weight
features.1.weight
features.1.bias
features.4.0.conv1.weight
features.4.0.bn1.weight
features.4.0.bn1.bias
features.4.0.conv2.weight
features.4.0.bn2.weight
features.4.0.bn2.bias
features.4.1.conv1.weight
features.4.1.bn1.weight
features.4.1.bn1.bias
features.4.1.conv2.weight
features.4.1.bn2.weight
features.4.1.bn2.bias
features.4.2.conv1.weight
features.4.2.bn1.weight
features.4.2.bn1.bias
features.4.2.conv2.weight
features.4.2.bn2.weight
features.4.2.bn2.bias
features.5.0.conv1.weight
features.5.0.bn1.weight
features.5.0.bn1.bias
features.5.0.conv2.weight
features.5.0.bn2.weight
features.5.0.bn2.bias
features.5.0.downsample.0.weight
features.5.0.downsample.1.weight
features.5.0.downsample.1.bias
features.5.1.conv1.weight
features.5.1.bn1.weight
features.5.1.bn1.bias
features.5.1.conv2.weight
features.5.1.bn2.weight
features.5.1.bn2.bias
features.5.2.conv1.weight
features.5.2.bn1.weight
features.5.2.bn1.bias
features.5.2.conv2.weight
features.5.2.bn2.weight
features.5.2.bn2.bias
features.5.3.conv1.weight
features.5.3.bn1.weight
features.5.3.bn1.bias
features.5.3.conv2.weight
features.5.3.bn2.weight
features.5.3.bn2.bias
features.6.0.conv1.weight
features.6.0.bn1.weight
features.6.0.bn1.bias
features.6.0.conv2.weight
features.6.0.bn2.weight
features.6.0.bn2.bias
features.6.0.downsample.0.weight
features.6.0.downsample.1.weight
features.6.0.downsample.1.bias
features.6.1.conv1.weight
features.6.1.bn1.weight
features.6.1.bn1.bias
features.6.1.conv2.weight
features.6.1.bn2.weight
features.6.1.bn2.bias
features.6.2.conv1.weight
features.6.2.bn1.weight
features.6.2.bn1.bias
features.6.2.conv2.weight
features.6.2.bn2.weight
features.6.2.bn2.bias
features.6.3.conv1.weight
features.6.3.bn1.weight
features.6.3.bn1.bias
features.6.3.conv2.weight
features.6.3.bn2.weight
features.6.3.bn2.bias
features.6.4.conv1.weight
features.6.4.bn1.weight
features.6.4.bn1.bias
features.6.4.conv2.weight
features.6.4.bn2.weight
features.6.4.bn2.bias
features.6.5.conv1.weight
features.6.5.bn1.weight
features.6.5.bn1.bias
features.6.5.conv2.weight
features.6.5.bn2.weight
features.6.5.bn2.bias
features.7.0.conv1.weight
features.7.0.bn1.weight
features.7.0.bn1.bias
features.7.0.conv2.weight
features.7.0.bn2.weight
features.7.0.bn2.bias
features.7.0.downsample.0.weight
features.7.0.downsample.1.weight
features.7.0.downsample.1.bias
features.7.1.conv1.weight
features.7.1.bn1.weight
features.7.1.bn1.bias
features.7.1.conv2.weight
features.7.1.bn2.weight
features.7.1.bn2.bias
features.7.2.conv1.weight
features.7.2.bn1.weight
features.7.2.bn1.bias
features.7.2.conv2.weight
features.7.2.bn2.weight
features.7.2.bn2.bias
classifier.0.weight
tensor([[-4.0103e-02, -3.0650e-03, -1.7690e-02,  1.8196e-02,  1.8940e-02,
         -5.2432e-02,  4.0119e-02,  1.6359e-02, -1.6112e-02,  3.6596e-02,
          1.1273e-02,  1.9971e-02,  2.4494e-02, -4.9495e-02, -2.4485e-02,
         -2.7787e-02, -4.4333e-02,  3.5293e-02, -1.1992e-02, -4.7242e-02,
          2.5577e-02, -3.0916e-02, -8.1142e-03, -2.4182e-03, -1.7207e-02,
          1.7690e-02,  3.6298e-02, -3.3379e-03,  1.5704e-02,  2.7359e-02,
         -1.6443e-04,  2.8785e-02, -4.1732e-03,  1.2785e-02,  6.6502e-02,
          2.3740e-02, -1.9601e-02,  2.0654e-02, -1.5341e-02,  5.1119e-02,
         -9.3292e-03, -4.6272e-02,  4.0541e-03, -1.6780e-02,  2.7696e-02,
         -2.8260e-02, -2.7696e-02, -2.6938e-02, -5.7942e-02, -1.7714e-02,
          1.1813e-02, -3.6892e-02,  3.2378e-02, -7.0157e-02,  2.9828e-03,
          4.5477e-02,  9.4537e-03,  1.4507e-02,  2.3831e-02,  1.3459e-02,
          5.2202e-02, -3.2785e-02, -1.7972e-02,  3.4436e-02, -1.7274e-02,
         -5.2432e-02,  3.5817e-02, -3.7003e-02,  3.9805e-02, -3.2142e-02,
          4.8118e-02,  3.0372e-02,  3.1019e-02,  2.0520e-02, -2.4969e-02,
          1.9197e-02, -5.7083e-02,  1.9971e-02,  4.1865e-03, -1.6170e-02,
         -2.7157e-02,  7.6236e-02,  7.7632e-04,  1.2858e-02,  6.0151e-02,
          4.4184e-02, -2.4115e-02,  3.8711e-03,  3.1806e-02, -7.6332e-03,
         -2.3273e-02, -3.6383e-02, -3.8611e-02,  4.8423e-02,  1.4127e-03,
          2.1603e-02,  5.7787e-03,  1.6463e-02,  2.6809e-02, -1.7455e-02,
         -3.9449e-03, -2.3069e-03,  6.2234e-02, -2.9278e-02, -9.0541e-03,
          4.0190e-02,  3.4797e-02, -3.7103e-03, -2.7790e-02, -4.0176e-02,
          4.5264e-03,  1.3271e-02, -3.0030e-03,  2.4259e-02, -2.2459e-02,
         -8.1706e-03,  2.9898e-02, -4.2108e-02, -3.8830e-03, -1.6584e-03,
         -2.5688e-03,  5.5268e-02,  2.5715e-03, -8.0469e-03, -5.3614e-03,
          3.2183e-02, -1.6894e-02,  4.4844e-02, -5.0618e-03,  3.0681e-03,
          3.4562e-02,  2.0505e-02,  5.2180e-02, -4.3831e-02,  4.5276e-02,
         -8.6777e-03,  4.2830e-02, -2.6109e-02, -4.2990e-03, -7.8125e-02,
          2.6773e-03, -4.4095e-03, -6.4931e-02, -1.2023e-02,  3.0119e-02,
          6.6624e-02,  3.2416e-02,  1.4566e-02, -6.7606e-02, -4.9717e-02,
         -2.2283e-02, -1.9339e-02, -8.6594e-02, -9.5567e-03, -3.6641e-02,
         -2.4269e-02,  6.0891e-03,  3.5978e-02, -1.3062e-02, -2.4939e-02,
         -2.6925e-02, -2.6563e-03, -1.3851e-02, -5.0669e-03,  6.0252e-02,
         -4.9383e-02,  4.5053e-02,  1.6863e-02,  2.6193e-02, -3.0816e-02,
          3.9407e-02, -5.0734e-02, -2.5574e-02,  1.5975e-02,  3.2121e-02,
         -2.1032e-02, -1.1800e-02,  2.0957e-02,  1.7678e-02,  3.2595e-02,
          2.7446e-02,  2.3684e-02,  4.6127e-03, -1.3901e-02, -1.4728e-02,
          5.2495e-02, -3.1069e-02,  1.1383e-02,  5.4902e-02, -9.8808e-03,
          5.5881e-02,  3.7348e-02,  4.9416e-03, -5.1515e-02, -6.1780e-02,
          2.5301e-02, -2.6315e-02, -1.8703e-02,  2.9316e-02, -2.3085e-02,
          4.1748e-03,  5.2638e-02, -2.1772e-02, -1.4431e-03,  9.2481e-03,
         -4.8283e-02,  3.1777e-03,  8.0173e-02,  4.5029e-03, -1.4431e-02,
         -3.3560e-02,  4.6829e-02, -9.9793e-03, -3.2405e-03, -1.1975e-02,
          5.5930e-02,  2.2951e-02, -3.9283e-02,  2.4410e-03,  3.8763e-03,
         -2.1583e-02, -4.4375e-03, -3.1224e-02,  2.5172e-02,  5.3837e-03,
         -4.2036e-02,  3.9306e-02, -1.9016e-02,  8.3942e-02,  1.7185e-02,
         -3.3008e-02,  1.2389e-02, -6.3560e-02,  2.6117e-02, -1.6864e-02,
          5.5821e-02, -1.2640e-03, -4.0874e-02, -1.3435e-02,  2.2319e-02,
         -3.9246e-02,  2.1261e-02, -2.4980e-02, -5.5888e-02, -5.4517e-03,
          2.5954e-02,  4.3477e-02,  5.3093e-02, -1.7187e-03,  5.0169e-02,
         -3.7164e-02,  2.7800e-02,  1.2954e-02,  3.6849e-03, -1.0594e-02,
         -1.0012e-02, -1.7303e-02, -2.1361e-02,  2.9390e-02, -2.0674e-02,
          3.5865e-02,  7.5633e-03,  2.3617e-02, -4.1382e-02,  2.9991e-02,
          3.0470e-04, -2.7890e-02, -1.8794e-02, -1.6211e-02,  3.3249e-02,
          1.9253e-02,  1.9588e-02,  2.2323e-02,  1.7697e-02, -7.8543e-03,
         -7.8463e-03,  2.2645e-02, -3.9645e-02, -3.3896e-02,  2.1476e-02,
          3.9840e-02, -7.6785e-03, -1.7353e-02, -5.0593e-02,  2.9839e-02,
          9.3777e-03, -5.1932e-05,  4.8221e-02,  1.5305e-02,  1.1562e-02,
          1.7175e-02,  1.9763e-02,  2.1498e-02, -6.4569e-03,  4.4490e-02,
         -2.1948e-02,  3.3371e-02, -2.6169e-02,  2.1575e-02,  5.0593e-02,
         -7.6582e-03,  7.3652e-02, -9.0516e-02,  3.0139e-03,  3.1726e-02,
         -5.5917e-02,  2.1198e-02, -2.0204e-02, -3.1504e-02,  8.3176e-03,
          6.8362e-02, -1.7631e-02,  3.2006e-02, -3.2219e-02, -2.3685e-02,
          2.5310e-02,  2.4241e-02,  1.3161e-02,  4.5711e-02,  6.0400e-02,
          3.2823e-02,  3.7996e-02, -6.6305e-03,  6.0794e-03,  3.6651e-02,
         -7.5762e-03,  2.3375e-03,  1.8069e-02, -3.9109e-02,  3.8477e-02,
         -2.9035e-02, -1.5453e-02, -6.7045e-03,  3.4121e-02, -9.8876e-03,
         -8.4192e-06,  5.2210e-03,  1.6493e-02, -6.2000e-02, -6.0531e-04,
          7.3830e-03, -4.0898e-05, -6.0627e-03, -2.8498e-02, -7.6559e-03,
          1.7166e-02, -1.6621e-02,  3.0013e-02, -2.9750e-02,  5.7559e-02,
         -1.9507e-02,  2.1800e-02, -6.5081e-02,  2.9435e-02, -9.1154e-03,
          1.4963e-02,  2.7812e-02,  3.7519e-02,  1.0546e-02, -5.9419e-02,
         -2.2688e-02,  2.9650e-02,  3.1546e-02,  2.7190e-03,  4.2752e-02,
         -9.8702e-03,  4.0616e-02, -8.7285e-03,  4.1171e-02, -2.8747e-03,
         -4.7521e-02,  1.4819e-02,  3.4308e-02, -2.3178e-02,  6.8566e-04,
          6.5807e-02, -1.5936e-02, -5.6867e-02, -2.0194e-02,  1.0089e-02,
          3.2515e-02,  2.7668e-02, -8.3925e-02, -4.6546e-02,  3.0311e-02,
         -4.4808e-02, -6.0378e-02,  2.9398e-02, -4.5278e-03, -3.3444e-02,
          1.7838e-02,  1.8011e-02,  1.1978e-02,  6.9284e-03, -4.7839e-02,
          6.3141e-03, -2.7778e-03, -6.6707e-02,  3.5588e-02,  3.3485e-02,
          4.3899e-02, -4.5265e-02,  6.8920e-03, -1.4840e-02,  2.2699e-02,
         -5.0180e-02, -5.5397e-02,  1.2932e-02, -1.6373e-02, -1.4470e-03,
         -3.1263e-02,  1.2203e-02,  6.5744e-02,  5.2748e-02,  4.9446e-02,
          1.1454e-02, -1.2506e-02,  1.6448e-02,  6.4812e-03,  9.1295e-03,
          4.7307e-02, -1.0583e-02,  2.6514e-02, -2.7043e-02,  2.1754e-02,
         -7.4832e-03,  8.8365e-03,  3.6107e-02,  1.0179e-02,  3.0531e-02,
          2.9151e-02,  2.7932e-02,  2.9156e-02,  2.0762e-02, -2.9590e-02,
         -8.0620e-04, -5.6288e-02,  4.2960e-02,  2.9789e-02,  1.9852e-02,
         -3.9940e-02, -1.6577e-03, -9.1024e-03,  1.2403e-03,  3.8063e-02,
          5.4522e-02,  3.3541e-02,  1.8009e-02, -2.9072e-03, -2.2482e-04,
         -6.2176e-03, -7.1625e-02, -2.9141e-02,  2.6403e-02, -3.7902e-02,
          7.6599e-03,  9.9294e-03,  4.3861e-02,  2.1696e-02,  3.6313e-03,
          5.7046e-02,  8.5943e-02, -1.0694e-02, -6.1576e-02,  3.4494e-02,
         -1.4768e-02, -1.5311e-03,  5.9259e-02,  2.4781e-02,  2.1675e-02,
          7.6399e-02, -9.0391e-03, -1.2334e-02, -2.4607e-02,  5.1892e-02,
         -2.1849e-02, -8.1642e-02,  2.2554e-02, -1.6767e-02,  2.3100e-02,
          1.6989e-02,  9.1284e-04,  1.9720e-02, -3.2603e-02,  3.9477e-03,
         -2.7493e-02,  2.6490e-02,  1.4810e-02, -9.9203e-02, -4.7352e-03,
          8.6435e-03,  7.8281e-02, -3.9165e-02,  3.1929e-02,  2.9405e-02,
          4.3515e-02,  4.6316e-02,  1.0432e-03,  3.8957e-02, -2.6859e-02,
         -4.3090e-03,  1.5592e-02,  5.6056e-02, -3.0031e-02, -7.8511e-03,
         -3.4687e-02,  1.6892e-02, -2.0491e-03, -1.4476e-02,  4.8247e-02,
         -1.3725e-02, -2.7866e-02]], device='cuda:0')
classifier.0.bias
</code></pre>
<h3>Multiclass code (gradient is none)</h3>
<pre><code class="lang-auto">original_model = models.__dict__["resnet34"](pretrained=True)
model = ResNet34_Model(original_model, 3)
model = model.cuda()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = 3e-4, weight_decay=0)
criterion = nn.CrossEntropyLoss(reduction='mean').cuda()

X_batch = X_batch.cuda() # X_batch and y_batch are pulled from the first iteration of a data loader
y_batch = y_batch.cuda()

# train for 1 epoch
model.train()
output = model(X_batch).squeeze()
output_var = torch.softmax(output, dim=1)    
_, output_class = torch.max(output_var, 1)
loss = torch.tensor(criterion(output_var, y_batch), requires_grad=True)
optimizer.zero_grad()
loss.backward()
optimizer.step()

# check that gradients exist
for n, p in model.named_parameters():
    print(n)
    if(p.requires_grad) and ("bias" not in n):
        print(p.grad)
</code></pre>
<h4>Output</h4>
<pre><code class="lang-auto">features.0.weight
features.1.weight
features.1.bias
features.4.0.conv1.weight
features.4.0.bn1.weight
features.4.0.bn1.bias
features.4.0.conv2.weight
features.4.0.bn2.weight
features.4.0.bn2.bias
features.4.1.conv1.weight
features.4.1.bn1.weight
features.4.1.bn1.bias
features.4.1.conv2.weight
features.4.1.bn2.weight
features.4.1.bn2.bias
features.4.2.conv1.weight
features.4.2.bn1.weight
features.4.2.bn1.bias
features.4.2.conv2.weight
features.4.2.bn2.weight
features.4.2.bn2.bias
features.5.0.conv1.weight
features.5.0.bn1.weight
features.5.0.bn1.bias
features.5.0.conv2.weight
features.5.0.bn2.weight
features.5.0.bn2.bias
features.5.0.downsample.0.weight
features.5.0.downsample.1.weight
features.5.0.downsample.1.bias
features.5.1.conv1.weight
features.5.1.bn1.weight
features.5.1.bn1.bias
features.5.1.conv2.weight
features.5.1.bn2.weight
features.5.1.bn2.bias
features.5.2.conv1.weight
features.5.2.bn1.weight
features.5.2.bn1.bias
features.5.2.conv2.weight
features.5.2.bn2.weight
features.5.2.bn2.bias
features.5.3.conv1.weight
features.5.3.bn1.weight
features.5.3.bn1.bias
features.5.3.conv2.weight
features.5.3.bn2.weight
features.5.3.bn2.bias
features.6.0.conv1.weight
features.6.0.bn1.weight
features.6.0.bn1.bias
features.6.0.conv2.weight
features.6.0.bn2.weight
features.6.0.bn2.bias
features.6.0.downsample.0.weight
features.6.0.downsample.1.weight
features.6.0.downsample.1.bias
features.6.1.conv1.weight
features.6.1.bn1.weight
features.6.1.bn1.bias
features.6.1.conv2.weight
features.6.1.bn2.weight
features.6.1.bn2.bias
features.6.2.conv1.weight
features.6.2.bn1.weight
features.6.2.bn1.bias
features.6.2.conv2.weight
features.6.2.bn2.weight
features.6.2.bn2.bias
features.6.3.conv1.weight
features.6.3.bn1.weight
features.6.3.bn1.bias
features.6.3.conv2.weight
features.6.3.bn2.weight
features.6.3.bn2.bias
features.6.4.conv1.weight
features.6.4.bn1.weight
features.6.4.bn1.bias
features.6.4.conv2.weight
features.6.4.bn2.weight
features.6.4.bn2.bias
features.6.5.conv1.weight
features.6.5.bn1.weight
features.6.5.bn1.bias
features.6.5.conv2.weight
features.6.5.bn2.weight
features.6.5.bn2.bias
features.7.0.conv1.weight
features.7.0.bn1.weight
features.7.0.bn1.bias
features.7.0.conv2.weight
features.7.0.bn2.weight
features.7.0.bn2.bias
features.7.0.downsample.0.weight
features.7.0.downsample.1.weight
features.7.0.downsample.1.bias
features.7.1.conv1.weight
features.7.1.bn1.weight
features.7.1.bn1.bias
features.7.1.conv2.weight
features.7.1.bn2.weight
features.7.1.bn2.bias
features.7.2.conv1.weight
features.7.2.bn1.weight
features.7.2.bn1.bias
features.7.2.conv2.weight
features.7.2.bn2.weight
features.7.2.bn2.bias
classifier.0.weight
None
classifier.0.bias
</code></pre>

ptrblck: <p>You are creating a new tensor as your loss, thus detaching the original output of your criterion from the computation graph:</p>
<pre><code class="lang-python">loss = torch.tensor(criterion(output_var, y_batch), requires_grad=True)
</code></pre>
<p>Instead just call:</p>
<pre><code class="lang-python">loss = criterion(output, y_batch)
loss.backward()
</code></pre>
<p>Note, that <code>nn.CrossEntropyLoss</code> expects raw logits as the model output, so don’t apply the softmax on your output tensor.</p>

ellax123: <p>Awesome, thanks for your help with this!</p>