Model seems to be learning fine, but performs badly when not doing a loss step

Omroth: <p>I’m baffled on this one.</p>
<p>I am training with this core loop (batch size is 1):</p>
<pre><code class="lang-auto">for image_data_s, demographics_s, actual_age_s in train_loader:
      image_data_s = image_data_s.to(device)
      demographics_s = demographics_s.to(device)
      actual_age_s = actual_age_s.to(device)

      predicted_ages = model(image_data_s, demographics_s)
      predicted_ages = predicted_ages.squeeze(1)

      X.append(actual_age_s.clone().detach().item())
      Y.append(predicted_ages.clone().detach().item())

      loss = criterion(predicted_ages, actual_age_s)
      
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
</code></pre>
<p>As you can see, I save the training prediction and target so I can plot them, to see how training is going.</p>
<p>By the second epoch, the results during training are already looking ok (and get much better by epoch 10):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c969a563312ab8534a9c30b28b6aa7aa019ed40.png" data-download-href="https://discuss.pytorch.org/uploads/default/0c969a563312ab8534a9c30b28b6aa7aa019ed40" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/c/0c969a563312ab8534a9c30b28b6aa7aa019ed40_2_668x500.png" alt="image" data-base62-sha1="1NmotHmHVMto3tCWfTWuHT6yJKU" width="668" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/c/0c969a563312ab8534a9c30b28b6aa7aa019ed40_2_668x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/c/0c969a563312ab8534a9c30b28b6aa7aa019ed40_2_1002x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c969a563312ab8534a9c30b28b6aa7aa019ed40.png 2x" data-dominant-color="FCF9F7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1244×931 57.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, when I try to use this trained model - just on the training data to be clear - <em>without</em> the backward step I get this output:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/e/0e172d781496400cebe9c77d679508773ff9fca8.png" data-download-href="https://discuss.pytorch.org/uploads/default/0e172d781496400cebe9c77d679508773ff9fca8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/e/0e172d781496400cebe9c77d679508773ff9fca8_2_672x500.png" alt="image" data-base62-sha1="20Elea8CkC0GfwGdHOZ3aVw8yDe" width="672" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/e/0e172d781496400cebe9c77d679508773ff9fca8_2_672x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/e/0e172d781496400cebe9c77d679508773ff9fca8_2_1008x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/e/0e172d781496400cebe9c77d679508773ff9fca8.png 2x" data-dominant-color="FCFAF8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1255×933 52.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I am completely stumped.  I get this by running the exact same code, directly after training, but by commenting out these lines:</p>
<pre><code class="lang-auto"> optimizer.zero_grad()
 loss.backward()
 optimizer.step()
</code></pre>
<p>What am I missing?</p>

ptrblck: <p>Could you post the model definition or (better) a minimal and executable code snippet to reproduce the issue, please?</p>

Omroth: <p>Hiya, thanks for the respone.  It’s going to be difficult to get snippet that runs, so let’s see if you can see anything off with the model def:</p>
<pre><code class="lang-auto">

n_input_channels = 1

n_latent_channels = 128

total_down_conv_channels = n_latent_channels

first_down_res = True 

first_down_res_stride = (2 if first_down_res else 1)

second_down_res = True

second_down_res_stride = (2 if second_down_res else 1)

#

kernel_size = 3
padding = 1

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        # [200, 256, 256, x]

        self.el0a = nn.Conv3d(in_channels=n_input_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=2, bias=False)   
        self.el0b = nn.ReLU(inplace=True)
        self.el0b_bn = nn.BatchNorm3d(total_down_conv_channels)
    
        # [100, 128, 128, 64]

        self.el1 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=first_down_res_stride, bias=False)   
        self.el2 = nn.ReLU(inplace=True)
        self.el2_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.max_pool_1 = nn.MaxPool3d(kernel_size=[2,2,2])

        # [50, 64, 64, 64]

        self.el3 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el4 = nn.ReLU(inplace=True)
        self.el4_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.el3a = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el4a = nn.ReLU(inplace=True)
        self.el4a_bn = nn.BatchNorm3d(total_down_conv_channels)


        self.mid_conv_dropout = nn.Dropout(p=0.3)

        self.max_pool_2 = nn.MaxPool3d(kernel_size=[2,2,2])

        # [50, 64, 64, 64]

        self.el5 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=second_down_res_stride, bias=False)   
        self.el6 = nn.ReLU(inplace=True)
        self.el6_bn = nn.BatchNorm3d(total_down_conv_channels)

        # [25, 32, 32, 64]

        self.el7 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el8 = nn.ReLU(inplace=True)
        self.el8_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.el7a = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el8a = nn.ReLU(inplace=True)
        self.el8a_bn = nn.BatchNorm3d(total_down_conv_channels)

        # [25, 32, 32, 64]

        self.el9 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=1, padding=0, stride=1, bias=False)

        self.max_pool_3 = nn.MaxPool3d(kernel_size=[2,2,2])

        self.el9_bn = nn.BatchNorm3d(total_down_conv_channels)

        #############

        self.global_pool = nn.AdaptiveAvgPool3d(output_size=1)

        self.fc1 = nn.Linear(n_latent_channels + demographics_length, 256).double()
        self.fc1_relu = nn.ReLU()

        self.mid_fc_dropout = nn.Dropout(p=0.3)

        self.fc2 = nn.Linear(256, 1).double()
        self.fc2_relu = nn.ReLU()

        #############

      

        
    def forward(self, image_data, demographics):
        x = self.el0a(image_data)
        x = self.el0b(x)
        x = self.el0b_bn(x)

        x = self.el1(x)
        x = self.el2(x)
        x = self.max_pool_1(x)
        x = self.el2_bn(x)
        x = self.el3(x)
        x = self.el4(x)
        x = self.max_pool_2(x)
        x = self.el4_bn(x)

        x = self.el3a(x)
        x = self.el4a(x)
        x = self.el4a_bn(x)

        x = self.mid_conv_dropout(x)

        x = self.el5(x)
        x = self.el6(x)
        x = self.el6_bn(x)

        x = self.el7(x)
        x = self.el8(x)
        x = self.el8_bn(x)

        x = self.el7a(x)
        x = self.el8a(x)
        x = self.el8a_bn(x)

        x = self.el9(x)
        x = self.max_pool_3(x)
        x = self.el9_bn(x)

        #


        x = self.global_pool(x) # this produces one per channel

        x = x.squeeze() # this is now shape [n_channels]
        x = torch.reshape(x, (1, total_down_conv_channels))

        x = torch.cat([x,demographics], axis=1)

        x = self.fc1(x)
        x = self.fc1_relu(x)

        x = self.mid_fc_dropout(x)

        x = self.fc2(x)
        x = self.fc2_relu(x)

        #

        predicted_age = x

        return predicted_age

</code></pre>

ptrblck: <p>Your code is missing a lot of definitions, so I won’t be able to execute and experiment with it.<br>
In any case, these lines look wrong:</p>
<pre><code class="lang-python">        x = x.squeeze() # this is now shape [n_channels]
        x = torch.reshape(x, (1, total_down_conv_channels))
</code></pre>
<p>as it seems you are forcing the batch size to be <code>1</code>.<br>
Is this always the case or could some unwanted broadcasting be executed somewhere?</p>

Omroth: <p>Hm.  I was under the mistaken impression that by the time data was going through Module::forward it was one sample.</p>
<p>(As it happens though, the batch size is 1 at all times)</p>
<p>If that line looks wrong though, we could try to fix it - would x.squeeze(1) be right?  The output of self.global_pool should be in shape [128,1] if I’m understanding AdaptiveAvgPool3d correctly, and I want it to be of shape [128] then concatinated with the demographics.</p>
<p>(I’ve added some definitions to the code above)</p>

Omroth: <p>I’ve removed the dropout layers for testing (and reduced the number of training samples to speed up by debugging, and this is interesting):</p>
<p>Epoch 2 during training:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/a/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6.png" data-download-href="https://discuss.pytorch.org/uploads/default/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6_2_669x499.png" alt="image" data-base62-sha1="vbhcRTCOXTX0n6O9PcfDcTbAYLA" width="669" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6_2_669x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6_2_1003x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/a/da89beeca1dcfdd8e1acabc96b3dfec0335e43a6.png 2x" data-dominant-color="FCFBFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1248×932 33 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Without training after Epoch 2 on same data:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/e/feb3aa62764deb618362d662b30d9a20b2b9960b.png" data-download-href="https://discuss.pytorch.org/uploads/default/feb3aa62764deb618362d662b30d9a20b2b9960b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feb3aa62764deb618362d662b30d9a20b2b9960b_2_672x500.png" alt="image" data-base62-sha1="AlccY8dmcTpsk4BWqcIThxDz38n" width="672" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feb3aa62764deb618362d662b30d9a20b2b9960b_2_672x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feb3aa62764deb618362d662b30d9a20b2b9960b_2_1008x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/e/feb3aa62764deb618362d662b30d9a20b2b9960b.png 2x" data-dominant-color="FCFCFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1241×923 30 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I bet this is something silly but I can’t see it</p>

ptrblck: <aside class="quote no-group" data-username="Omroth" data-post="6" data-topic="173782">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/o/fbc32d/48.png" class="avatar"> Omroth:</div>
<blockquote>
<p>Without training after Epoch 2 on same data:</p>
</blockquote>
</aside>
<p>Did you keep the model in <code>.train()</code> mode? If so, note that the batchnorm stats would still be updated. However, if you never switch to <code>.eval()</code> these running stats also won’t be used, so it’s still unclear to me what’s causing the issue.<br>
Could you describe your workflow again so that I could try to use random data to reproduce it?<br>
I.e. in particular when <code>.train/eval</code> is called, when the model is updated etc.</p>

Omroth: <p>I do keep the model in train mode (I’ve tried going to eval first, but that didn’t work.)</p>
<p>I’ve created a version which should work on its own without any edits - thank you for taking a look.</p>
<pre><code class="lang-auto">

import torch
from torch import nn
import numpy as np
from torch import nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, TensorDataset
from matplotlib import pyplot as plt
import numpy as np

###############

saved_model_filename = 'brain_age_network.pt'

###

learning_rate = 1e-3

batch_size = 1

n_global_epochs = 2 # going around all of the scans


###

rescaled_image_resolution = [200,256,256]

#

dropout = False

n_input_channels = 1

n_latent_channels = 128

total_down_conv_channels = n_latent_channels

first_down_res = True 

first_down_res_stride = (2 if first_down_res else 1)

second_down_res = True

second_down_res_stride = (2 if second_down_res else 1)

#

kernel_size = 3
padding = 1


###

class RandDataset(Dataset):
    def __init__(self, seed):
        np.random.seed(seed)

        self.image_data = np.random.rand(1, rescaled_image_resolution[0],rescaled_image_resolution[1],rescaled_image_resolution[2])

        self.age = np.random.random()

    def __len__(self):
        return 1

    def __getitem__(self, idx):
        return self.image_data, self.age


#

class BNModel(nn.Module):
    def __init__(self):
        super().__init__()

        # [200, 256, 256, x]

        self.el0a = nn.Conv3d(in_channels=n_input_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=2, bias=False)   
        self.el0b = nn.ReLU(inplace=True)
        self.el0b_bn = nn.BatchNorm3d(total_down_conv_channels)
    
        # [100, 128, 128, 64]

        self.el1 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=first_down_res_stride, bias=False)   
        self.el2 = nn.ReLU(inplace=True)
        self.el2_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.max_pool_1 = nn.MaxPool3d(kernel_size=[2,2,2])

        # [50, 64, 64, 64]

        self.el3 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el4 = nn.ReLU(inplace=True)
        self.el4_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.el3a = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el4a = nn.ReLU(inplace=True)
        self.el4a_bn = nn.BatchNorm3d(total_down_conv_channels)

        if dropout:
            self.mid_conv_dropout = nn.Dropout(p=0.3)

        self.max_pool_2 = nn.MaxPool3d(kernel_size=[2,2,2])

        # [50, 64, 64, 64]

        self.el5 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=second_down_res_stride, bias=False)   
        self.el6 = nn.ReLU(inplace=True)
        self.el6_bn = nn.BatchNorm3d(total_down_conv_channels)

        # [25, 32, 32, 64]

        self.el7 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el8 = nn.ReLU(inplace=True)
        self.el8_bn = nn.BatchNorm3d(total_down_conv_channels)

        self.el7a = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=kernel_size, padding=padding, stride=1, bias=False)
        self.el8a = nn.ReLU(inplace=True)
        self.el8a_bn = nn.BatchNorm3d(total_down_conv_channels)

        # [25, 32, 32, 64]

        self.el9 = nn.Conv3d(in_channels=total_down_conv_channels,out_channels=total_down_conv_channels,kernel_size=1, padding=0, stride=1, bias=False)

        self.max_pool_3 = nn.MaxPool3d(kernel_size=[2,2,2])

        self.el9_bn = nn.BatchNorm3d(total_down_conv_channels)

        #############

        self.global_pool = nn.AdaptiveAvgPool3d(output_size=1)

        self.fc1 = nn.Linear(n_latent_channels, 256)

        self.fc1_relu = nn.ReLU()

        if dropout:
            self.mid_fc_dropout = nn.Dropout(p=0.3)

        self.fc2 = nn.Linear(256, 1)

        self.fc2_relu = nn.ReLU()

        #############

      

        
    def forward(self, image_data):
        x = self.el0a(image_data)
        x = self.el0b(x)
        x = self.el0b_bn(x)

        x = self.el1(x)
        x = self.el2(x)
        x = self.max_pool_1(x)
        x = self.el2_bn(x)
        x = self.el3(x)
        x = self.el4(x)
        x = self.max_pool_2(x)
        x = self.el4_bn(x)

        x = self.el3a(x)
        x = self.el4a(x)
        x = self.el4a_bn(x)

        if dropout:
            x = self.mid_conv_dropout(x)

        x = self.el5(x)
        x = self.el6(x)
        x = self.el6_bn(x)

        x = self.el7(x)
        x = self.el8(x)
        x = self.el8_bn(x)

        x = self.el7a(x)
        x = self.el8a(x)
        x = self.el8a_bn(x)

        x = self.el9(x)
        x = self.max_pool_3(x)
        x = self.el9_bn(x)

        #

        x = self.global_pool(x) # this produces one per channel

        x = torch.reshape(x, (1, total_down_conv_channels))
        

        x = self.fc1(x)
        x = self.fc1_relu(x)

        if dropout:
            x = self.mid_fc_dropout(x)

        x = self.fc2(x)
        x = self.fc2_relu(x)

        #

        predicted_age = x

        return predicted_age


###



def train_model():
    # set random seeds
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)

    # setup device cuda vs. cpu
    cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if cuda else "cpu")

    assert cuda, "we require cuda for the brain network"

    model = BNModel().to(device)

    #

    training_seeds = [np.random.randint(100000,999999) for i in range(50)]

    # Setting the optimiser
    
    try:
        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=learning_rate,
            fused=True, # this does more of the work on the GPU
        )
    except Exception as e:
        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=learning_rate,
        )

    #

    criterion = torch.nn.MSELoss(reduction='sum')
    
    X = []
    Y = []

   #

    from matplotlib.backends.backend_pdf import PdfPages

    pdf = PdfPages("brain_age_model_training.pdf")

    #


    for global_epoch in range(1, n_global_epochs+1):
        global_epoch_train_loss = 0
        n_global_epoch_trainings = 0

        X = []
        Y = []

        for seed in training_seeds:
            try:
                dataset = RandDataset(seed)
            except Exception as e:
                continue

            train_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)

            # train for one epoch
            model.train()
            
            for image_data_s, actual_age_s in train_loader:
                image_data_s = image_data_s.float().to(device)
                actual_age_s = actual_age_s.float().to(device)

                

                # ===================forward=====================
                predicted_ages = model(image_data_s)

                predicted_ages = predicted_ages.squeeze(1)

                X.append(actual_age_s.clone().detach().item())
                Y.append(predicted_ages.clone().detach().item())

                loss = criterion(predicted_ages, actual_age_s)
                
                local_loss = loss.clone().detach()
                if global_epoch_train_loss is None: # keep it on the gpu to avoid unneccesary cpu/gpu syncs
                    global_epoch_train_loss = local_loss
                else:
                    global_epoch_train_loss += local_loss
                n_global_epoch_trainings += 1


                # ===================backward====================
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

          
    

        global_epoch_train_loss = global_epoch_train_loss.item()
        average_loss_per_scan = global_epoch_train_loss / n_global_epoch_trainings
        print(f"Global epoch average scan loss: {average_loss_per_scan}")

        # save a graph to a pdf to show our predictions vs reality

        fig, ax = plt.subplots()
        ax.plot([0.5,1.0],[0.5,1.0],c="tab:grey")
        ax.scatter(X, Y, s=2.2,c="tab:orange")
        ax.set_xlabel("Actual age")
        ax.set_ylabel("Brain age")
        fig.suptitle(f"Epoch {global_epoch}.  Average loss per subject: {average_loss_per_scan}", fontsize=8)
        pdf.savefig(fig)
        plt.close(fig)




    if True:
        X = []
        Y = []

        #with torch.no_grad():
        if True:
            #model.train()
            #model.eval()

            for seed in training_seeds:
                try:
                    dataset = RandDataset(seed)
                except Exception as e:
                    continue

                train_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)

                for image_data_s, actual_age_s in train_loader:

                    image_data_s = image_data_s.float().to(device)
                    actual_age_s = actual_age_s.float().to(device)

                    # ===================forward=====================
                    predicted_ages = model(image_data_s)
                    predicted_ages = predicted_ages.squeeze(1)

                    X.append(actual_age_s.clone().detach().item())
                    Y.append(predicted_ages.clone().detach().item())

                    loss = criterion(predicted_ages, actual_age_s)


        #

        fig, ax = plt.subplots()
        ax.plot([0.5,1.0],[0.5,1.0],c="tab:grey")
        ax.scatter(X, Y, s=2.2,c="tab:orange")
        ax.set_xlabel("Actual age")
        ax.set_ylabel("Brain age")
        fig.suptitle(f"POST TRAINING RESULTS NON-RELOADED SCANS", fontsize=8)
        pdf.savefig(fig)
        plt.close(fig)
        
    

    #

    # save the state dict of the model
    torch.save(model.state_dict(), saved_model_filename)

    #

    pdf.close()
    plt.close('all')


###


train_model()

</code></pre>
<p>This gives the same fundamental problem to me - during training we get at least some learning (although it’s more random with fully random data):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/f/bf931d078caa6be308334389684abc7817f6db2b.png" data-download-href="https://discuss.pytorch.org/uploads/default/bf931d078caa6be308334389684abc7817f6db2b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf931d078caa6be308334389684abc7817f6db2b_2_690x499.png" alt="image" data-base62-sha1="rkKnVdOCv8QeomJM0we2WJpMsjN" width="690" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf931d078caa6be308334389684abc7817f6db2b_2_690x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf931d078caa6be308334389684abc7817f6db2b_2_1035x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/f/bf931d078caa6be308334389684abc7817f6db2b.png 2x" data-dominant-color="FCFCFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1273×922 30.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>And yet when not doing the backwards step (the final graph) - we just get flat output:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/2/d242190cf5fe992c1a58c21e5f8097e54249c658.png" data-download-href="https://discuss.pytorch.org/uploads/default/d242190cf5fe992c1a58c21e5f8097e54249c658" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d242190cf5fe992c1a58c21e5f8097e54249c658_2_681x500.png" alt="image" data-base62-sha1="u01SJso9Nll7egohOYMFg5T8jJm" width="681" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d242190cf5fe992c1a58c21e5f8097e54249c658_2_681x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/2/d242190cf5fe992c1a58c21e5f8097e54249c658_2_1021x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/2/d242190cf5fe992c1a58c21e5f8097e54249c658.png 2x" data-dominant-color="FCFCFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1261×925 30.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

Omroth: <p>Any chance you could take a look at this <a class="mention" href="/u/ptrblck">@ptrblck</a> ?  I’m still stumped.</p>

ptrblck: <p>I think the lack of optimizing the parameters in the second run lets the model collapse.<br>
If I remove the <code>optimizer.step()</code> call from the first loop, I see:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d8964177557d900729f5fa5913016566b2a7eb8f.png" alt="image" data-base62-sha1="uU13pwDWhGeHbUhB74JJdE8mG9F" width="386" height="291"></p>
<p>Additionally, the last <code>nn.ReLU</code> layer might also cause trouble as it would clip the output and could show the <code>bias</code> only.</p>

Omroth: <p>Thank you for taking a look.</p>
<p>That final relu does indeed look like a mistake, and I’ve removed it.</p>
<p>But I don’t understand the idea of “model collapse”.  All that second loop is doing is testing the model, not training it - why is it possible for it to have a constant output?</p>

ptrblck: <p>You could check if only the bias of the last layer is used and all previous activations are clipped to zero.<br>
Also, this effect is not specific to the second run as it’s already visible in the first loop if no parameters are updated.</p>

Omroth: <p>I think the culprit might be the AdaptiveAvgPool3d layer.  Is there an alternative just for testing to get the [25, 32, 32, 64] shape 3D data into a shortish linear array?</p>

ptrblck: <p>Based on the definition of <code>global_pool = nn.AdaptiveAvgPool3d(output_size=1)</code> it seems you want to reduce the depth, height, and width to <code>1</code> via a <code>mean</code> operation.<br>
An alternative would be to use <code>max</code> or <code>sum</code>, but I don’t know if this would help or if increasing the output size might be beneficial in your case.</p>

Omroth: <p>Well, increasing the output size of the <code>AdaptiveAvgPool3d</code> layer has stopped the output being flat, but it still isn’t the same as it is during the training runs.  I guess I think the problem is in the reshape call, but I don’t know how else I’m supposed to get data which is in shape, say</p>
<p>[1,128,4,4,4]</p>
<p>into</p>
<p>[1,128 * 4 * 4 * 4]</p>
<p>So that it can be put through a linear layer.</p>
<p>And I still have no idea why the behavour <strong>changes</strong> when you aren’t doing a backwards step.  Do yo understand why?</p>

ptrblck: <p>I guess because your model is always saturating the outputs which explains also the static output during training if no parameters are updated.<br>
If you call <code>optimizer.step()</code> in each iteration the output will change slightly and the training loop looks as if the outputs are improving while you might just have shifted the static output value. You could test it by updating once and checking if the next outputs would all have the same but new value again.<br>
In the end I don’t think your model is training properly as not even the training samples can overfit the target, so I would recommend focusing on this issue first.</p>

Omroth: <p>Thank you as always for your help, I think that gets to the heart of the matter.</p>