Model not training: Model predicting same thing every time

maycaj: <p>Basically, my objective was to remove the classifying layer from the model (commented out) and have the output of the second convolutional layer act as a classification of each group of pixels.</p>
<p>Here is my model definition:</p>
<pre><code class="lang-auto">class TinyVGG(nn.Module):
    """
    Model architecture copying TinyVGG from: 
    https://poloclub.github.io/cnn-explainer/
    """
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -&gt; None:
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, 
                      out_channels=hidden_units, 
                      kernel_size=3, 
                      stride=1, 
                      padding=1), 
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, 
                      out_channels=hidden_units,
                      kernel_size=3,
                      stride=1,
                      padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,
                         stride=2) # default stride value is same as kernel_size
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_units, len(train_data_custom.classes), kernel_size=3, padding=1), # modified to add 3 for 3 different classes
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # self.classifier = nn.Sequential(
        #     nn.Flatten(),
        #     # Where did this in_features shape come from? 
        #     # It's because each layer of our network compresses and changes the shape of our inputs data.
        #     # nn.Linear(in_features= 31360, # modified from: hidden_units*16*16,
        #     #           # above is same as 10*130*174
        #     #           out_features=output_shape)
        #     nn.LazyLinear(out_features=output_shape)
        # )

    def forward(self, x: torch.Tensor):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        # x = self.classifier(x)
        return x

model_0 = TinyVGG(input_shape=128, # number of color channels (3 for RGB) - modified
                  hidden_units=10, 
                  output_shape=len(train_data_custom.classes)).to(device)
</code></pre>
<p>Here is my training step:</p>
<pre><code class="lang-auto">def train_step(model: torch.nn.Module, 
               dataloader: torch.utils.data.DataLoader, 
               loss_fn: torch.nn.Module, 
               optimizer: torch.optim.Optimizer,
               epoch):
    # Put model in train mode - as apposed to evaluation mode - model.eval()
    model.train()
    
    # Setup train loss and train accuracy values
    train_loss, train_acc = 0, 0
    
    # Loop through data loader data batches
    for batch, (X, y) in enumerate(dataloader):
        # Send data to target device
        X = X.to(device)
        print(f"[train_step] X size: {X.size()}")
        # 1. Forward pass
        y_pred = model(X) 
        print(f"[train_step] y_pred size: {y_pred.size()}")
        y = y.to(device)
        y = y.long()
        # 2. Calculate  and accumulate loss
        y = torch.tile(y.view((-1, 1, 1)), (1, 174, 130))
        print(f"[train_step] y size: {y.size()}")
        loss = loss_fn(y_pred, y)
        # print(f"[train_step] loss: {loss}")

        train_loss += loss.mean().item() # modified
        
        print(f"[train_step] train_loss: {train_loss/(batch+1)}")
        # 3. Optimizer zero grad - ensure loss.backward() isn't accumulating the gradient from previous steps
        optimizer.zero_grad()

        # 4. Loss backward - determine what direction we will step in
        loss = loss.mean()
        loss.backward()


        # Log gradients
        for name, param in model.named_parameters():
            writer.add_histogram(f'gradients/{name}', param.grad, epoch)
        writer.add_histogram(f'train_loss/', train_loss/(batch+1), epoch)
        writer.add_histogram(f'X', X, epoch)

        # 5. Optimizer step
        optimizer.step()

        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)

        train_acc += (y_pred_class == y).sum().item()/(len(y_pred)*y_pred_class.size(dim=1)*y_pred_class.size(dim=2))
        print(f"[train_step] Train_acc: {train_acc/(batch+1)}")


    # Adjust metrics to get average loss and accuracy per batch 
    train_loss = train_loss / len(dataloader)
    train_acc = train_acc / len(dataloader)
    return train_loss, train_acc
</code></pre>
<p>Here are the shapes going into and out of the model:</p>
<pre><code class="lang-auto">[train_step] X size: torch.Size([9, 128, 696, 520])
[train_step] y_pred size: torch.Size([9, 3, 174, 130])
[train_step] y size: torch.Size([9, 174, 130])
</code></pre>
<p>The loss decreases from epoch 1 to epoch 2, but the testing accuracy is at chance because the model predicts the second category for every image.</p>
<pre><code class="lang-auto">Epoch: 1 | train_loss: 111.7413 | train_acc: 0.4336 | test_loss: 1.1037 | test_acc: 0.3333
Epoch: 2 | train_loss: 1.1082 | train_acc: 0.3333 | test_loss: 1.1094 | test_acc: 0.3333
Epoch: 3 | train_loss: 1.1099 | train_acc: 0.3333 | test_loss: 1.1051 | test_acc: 0.3333
</code></pre>

J_Johnson: <p>Have you applied rescaling weights to the loss function?</p>

maycaj: <p>From searching around a bit, I’m not sure how to do that. Can you link a resource so I can look into it?</p>

J_Johnson: <p>This is a very common problem. I wrote something in the FAQ which addresses it with a code example.</p>
<aside class="quote quote-modified" data-post="1" data-topic="194579">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/j_johnson/48/55494_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-do-i-get-the-accuracy-of-an-unbalanced-dataset-or-segmentation-task/194579">How do I get the accuracy of an unbalanced dataset or segmentation task?</a> <a class="badge-category__wrapper " href="/c/faq/30"><span data-category-id="30" style="--category-badge-color: #0088CC; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="The FAQ category contains commonly-asked questions and their answers. Please refer to this section before you post your query."><span class="badge-category__name">FAQ</span></span></a>
  </div>
  <blockquote>
    This post is to define a Class Weighted Accuracy function(WCA). As questions related to this get asked often, I thought it might help people to post a tool torchers can use and reference here. 
Why? 
Take, for example, a classification dataset of kittens and puppies with a ratio of 0.2:0.8 kittens to puppies. A model trained on this dataset might show an overall accuracy of 80% by just predicting “puppies” every time. 
A WCA function resolves this problem by checking the accuracy of each class s…
  </blockquote>
</aside>

<p>In a nut shell, when you train a supervised task on an unbalanced dataset, your model will optimize for decreasing loss, which can mean only answering the option that is the most common in the dataset.</p>
<p>To address this, we need to reweight those class indices and pass that set of weights into the loss function.</p>
<p>EDIT: I just reviewed the FAQ in the link and it seems that only addresses assessing an unbalanced dataset for model accuracy. I had later intended to include a function within it to calculate the rescaling weights, but it seems I never got around to that. I will update it in the next 24 hours to include that.</p>

J_Johnson: <p>I’ve added a reply to the FAQ which gives a function for calculating the rescaling weight vector used by the CrossEntropyLoss function in PyTorch.</p>
<p>You can find that here:</p>
<aside class="quote quote-modified" data-post="2" data-topic="194579">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/j_johnson/48/55494_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-do-i-get-the-accuracy-of-an-unbalanced-dataset-or-segmentation-task/194579/2">How do I get the accuracy of an unbalanced dataset or segmentation task?</a> <a class="badge-category__wrapper " href="/c/faq/30"><span data-category-id="30" style="--category-badge-color: #0088CC; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="The FAQ category contains commonly-asked questions and their answers. Please refer to this section before you post your query."><span class="badge-category__name">FAQ</span></span></a>
  </div>
  <blockquote>
    Here is a function for calculating the weight vector used by CrossEntropyLoss based on a set of targets with integers. 
def scaled_weights(targets):
    '''
    Gets class weights
    :param targets: pass in the matrix of size (batch, classes) containing correct targets or
        labels as one_hot
    :return: returns a vector for weight rescaling for use in cross entropy loss function
    '''

    N, C = targets.size(0), targets.size(1)

    # get vector of total for each class
    totals = ta…
  </blockquote>
</aside>