Model cannot learn, traing loss is not converge

username: quyujunbbb
cooked: <p>Hi, I am trying to train a single non-local block (as described in <a href="https://arxiv.org/abs/1711.07971" class="inline-onebox" rel="noopener nofollow ugc">[1711.07971] Non-local Neural Networks</a>), but the network seems not converge at all.</p>
<p>The general idea is (i) taking a batch of extracted features (bs, c, t, h, w) from customized video clips as input (I use ResNet 3D as the backbone), (ii) forwarding the input into a non-local block, (iii) output a score between [0, 1] as a regression task.</p>
<p>I trained the model with 5-fold cross validation, 60 epochs, batch size 16, initial learning rate 0.001 and reduced by 0.1 every 20 epochs, loss function MSE, optimizer Adam.</p>
<p>I got the averaged loss on test set being 0.047283. However, the plotted loss figures seem indicating the training does not learn anything… The training loss doesn’t decrease at all. I tried to change the batch size to 32, learning rate start from 0.0001, the results are very similar.</p>
<p>I am not sure if my codes have errors or any steps got wrong. <img src="https://discuss.pytorch.org/images/emoji/apple/cry.png?v=12" title=":cry:" class="emoji" alt=":cry:" loading="lazy" width="20" height="20"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/c/fced605c6cf39038031852fb565adfecfff3f08b.png" data-download-href="https://discuss.pytorch.org/uploads/default/fced605c6cf39038031852fb565adfecfff3f08b" title="losses"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fced605c6cf39038031852fb565adfecfff3f08b_2_267x500.png" alt="losses" data-base62-sha1="A5uTKoBlpN9rHoMihhuycbxeg0b" width="267" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fced605c6cf39038031852fb565adfecfff3f08b_2_267x500.png, https://discuss.pytorch.org/uploads/default/original/3X/f/c/fced605c6cf39038031852fb565adfecfff3f08b.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/f/c/fced605c6cf39038031852fb565adfecfff3f08b.png 2x" data-dominant-color="FAF7F6"><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">losses</span><span class="informations">373×698 42.3 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>
<pre data-code-wrap="python"><code class="lang-python"># training code
def train(data):
    k_folds = 5
    results = {}
    kfold = KFold(n_splits=k_folds, shuffle=True)
    loss_func = nn.MSELoss()

    for fold, (train_ids, test_ids) in enumerate(kfold.split(data)):
        train_subsampler = SubsetRandomSampler(train_ids)
        test_subsampler = SubsetRandomSampler(test_ids)
        train_loader = DataLoader(data,
                                  batch_size=16,
                                  sampler=train_subsampler)
        test_loader = DataLoader(data,
                                 batch_size=16,
                                 sampler=test_subsampler)

        net = Nonlocal()
        if torch.cuda.is_available():
            net = nn.DataParallel(net)
            net.cuda()

        opt = torch.optim.Adam(net.parameters(), lr=0.001)
        scheduler = StepLR(opt, step_size=20, gamma=0.1)

        for epoch_index in range(60):
            torch.set_grad_enabled(True)
            net.train()
            for train_batch_index, (img_batch,
                                    label_batch) in enumerate(train_loader):
                if torch.cuda.is_available():
                    img_batch = img_batch.cuda().float()
                    label_batch = label_batch.cuda().float()

                predict = net(img_batch)
                loss = loss_func(predict, label_batch)

                net.zero_grad()
                loss.backward()
                opt.step()

            torch.set_grad_enabled(False)
            net.eval()
            total_loss = []
            total_sample = 0

            for test_batch_index, (img_batch,
                                   label_batch) in enumerate(test_loader):
                if torch.cuda.is_available():
                    img_batch = img_batch.cuda().float()
                    label_batch = label_batch.cuda().float()

                predict = net(img_batch)
                loss = loss_func(predict, label_batch)

                total_loss.append(loss)
                total_sample += img_batch.size(0)

            net.train()

            mean_loss = sum(total_loss) / total_loss.__len__()
            results[fold] = mean_loss
            print(f'(Test) loss:{mean_loss.item():4f}')

            scheduler.step()

    fold_loss = 0.0
    for key, value in results.items():
        fold_loss += value

# Network
class Nonlocal(nn.Module):

    def __init__(self):
        super(Nonlocal_FC, self).__init__()

        self.nl = NONLocalBlock3D(in_channels=1024)
        self.maxpool = nn.MaxPool3d(kernel_size=(4, 7, 7))
        self.avgpool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.fc = nn.Sequential(nn.Linear(in_features=1024, out_features=1))

    def forward(self, x):
        batch_size = x.size(0)
        x = self.nl(x)
        x = self.avgpool(x).view(batch_size, -1)
        x = self.fc(x)
        return x

# 3D Non-local block
# The code from https://github.com/AlexHex7/Non-local_pytorch/blob/master/lib/non_local_gaussian.py is used.
</code></pre>