Sudden big loss in a branch makes losses in other branches not optimized any more

username: Flicic_Suo
cooked: <p>I’m training an object detection model with YOLO structure:</p>
<pre><code class="lang-auto">class YOLO(nn.Module):
    def __init__(self, ...)
        self.backbone = resnet50(pretrained=True)  # pretrained weights, classification layer already removed
        self.neck = nn.Sequential(
            conv3x3(in_planes=2048, out_planes=2048),
            nn.ReLU(),
            conv3x3(in_planes=2048, out_planes=2048),
            nn.ReLU(),
            conv3x3(in_planes=2048, out_planes=2048),,
            nn.ReLU(),
        )
        self.heads = nn.ModuleList(
            [YOLOHead(num_classes) for anchor in self.anchors]
        )
    def forward(self, x, ...):
        x = self.backbone(x)
        x = self.neck(x)
        outs = [head(x) for head in self.heads]

class YOLOHead(nn.Module):
    def __init__(self, num_classes):
        super(YOLOHead, self).__init__()
        self.objectness = conv1x1(in_planes=2048, out_planes=1)
        self.cls = conv1x1(in_planes=2048, out_planes=1+num_classes)
        self.reg = conv1x1(in_planes=2048, out_planes=4)

    def forward(self, x):
        objectness = self.objectness(x)
        cls = self.cls(x)
        reg = self.reg(x)
        return objectness, cls, reg
</code></pre>
<p>I optimize the sum of loss of (1)objectness, (2)classification and (3)bbox regression. As shown above, the heads compute them separately, yet with a shared backbone.<br>
The issue is that, sometimes regression loss (mse_loss between gt and pred) may get a sudden raise, and then optimzed down to a moderate level again. But after that optimzing-down, loss of objectness and classification can hardly be optimized any more. My words may be vague but it looks like this (see the stable obj_loss and cls_loss after batch 26):</p>
<pre><code class="lang-auto">[epoch: 1, batch: 1] loss: 7.862009, obj_loss:  1.040800, cls_loss:  4.567673, reg_loss:  2.253536
[epoch: 1, batch: 2] loss: 7.364972, obj_loss:  1.027697, cls_loss:  4.538699, reg_loss:  1.798576
[epoch: 1, batch: 3] loss: 7.077137, obj_loss:  1.006557, cls_loss:  4.509902, reg_loss:  1.560679
[epoch: 1, batch: 4] loss: 6.991296, obj_loss:  0.959222, cls_loss:  4.449108, reg_loss:  1.582966
[epoch: 1, batch: 5] loss: 7.313764, obj_loss:  0.869381, cls_loss:  4.270252, reg_loss:  2.174131
[epoch: 1, batch: 6] loss: 6.717353, obj_loss:  0.832326, cls_loss:  4.162328, reg_loss:  1.722699
[epoch: 1, batch: 7] loss: 7.118721, obj_loss:  0.969238, cls_loss:  3.606729, reg_loss:  2.542755
[epoch: 1, batch: 8] loss: 6.670866, obj_loss:  0.907451, cls_loss:  3.944198, reg_loss:  1.819217
[epoch: 1, batch: 9] loss: 7.152078, obj_loss:  0.937133, cls_loss:  4.139555, reg_loss:  2.075390
[epoch: 1, batch: 10] loss: 6.950064, obj_loss:  0.961618, cls_loss:  4.229391, reg_loss:  1.759055
[epoch: 1, batch: 11] loss: 6.504709, obj_loss:  0.969900, cls_loss:  4.112210, reg_loss:  1.422600
[epoch: 1, batch: 12] loss: 6.989936, obj_loss:  0.993160, cls_loss:  3.952706, reg_loss:  2.044071
[epoch: 1, batch: 13] loss: 6.952768, obj_loss:  0.883456, cls_loss:  3.616553, reg_loss:  2.452760
[epoch: 1, batch: 14] loss: 7.013877, obj_loss:  0.963605, cls_loss:  4.107985, reg_loss:  1.942287
[epoch: 1, batch: 15] loss: 6.275753, obj_loss:  0.981065, cls_loss:  3.834120, reg_loss:  1.460567
[epoch: 1, batch: 16] loss: 6.559465, obj_loss:  0.931920, cls_loss:  3.854708, reg_loss:  1.772837
[epoch: 1, batch: 17] loss: 6.501263, obj_loss:  0.988465, cls_loss:  4.012724, reg_loss:  1.500074
[epoch: 1, batch: 18] loss: 6.336897, obj_loss:  0.972874, cls_loss:  3.923963, reg_loss:  1.440061
[epoch: 1, batch: 19] loss: 7.756456, obj_loss:  0.907921, cls_loss:  4.357026, reg_loss:  2.491509
[epoch: 1, batch: 20] loss: 6.757920, obj_loss:  1.009806, cls_loss:  4.316257, reg_loss:  1.431857
[epoch: 1, batch: 21] loss: 6.218307, obj_loss:  0.975769, cls_loss:  4.141428, reg_loss:  1.101111
[epoch: 1, batch: 22] loss: 11.127349, obj_loss:  1.423922, cls_loss:  6.209508, reg_loss:  3.493918
[epoch: 1, batch: 23] loss: 7.093385, obj_loss:  1.035403, cls_loss:  4.525311, reg_loss:  1.532671
[epoch: 1, batch: 24] loss: 7.644183, obj_loss:  1.034627, cls_loss:  4.520189, reg_loss:  2.089367
[epoch: 1, batch: 25] loss: 6.772291, obj_loss:  1.016327, cls_loss:  4.244712, reg_loss:  1.511253
[epoch: 1, batch: 26] loss: 193.881104, obj_loss:  1.830120, cls_loss:  17.074070, reg_loss:  174.976913
[epoch: 1, batch: 27] loss: 7.431651, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.825148
[epoch: 1, batch: 28] loss: 7.192277, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.585774
[epoch: 1, batch: 29] loss: 6.893618, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.287114
[epoch: 1, batch: 30] loss: 7.638692, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  2.032189
[epoch: 1, batch: 31] loss: 7.449893, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.843388
[epoch: 1, batch: 32] loss: 7.655496, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.048993
[epoch: 1, batch: 33] loss: 7.145531, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.539027
[epoch: 1, batch: 34] loss: 7.390171, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.783668
[epoch: 1, batch: 35] loss: 7.374976, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.768474
[epoch: 1, batch: 36] loss: 7.157423, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.550920
[epoch: 1, batch: 37] loss: 7.023059, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.416556
[epoch: 1, batch: 38] loss: 7.025788, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.419284
[epoch: 1, batch: 39] loss: 7.916700, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.310196
[epoch: 1, batch: 40] loss: 7.503070, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.896566
[epoch: 1, batch: 41] loss: 7.642205, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.035701
[epoch: 1, batch: 42] loss: 7.002790, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.396288
[epoch: 1, batch: 43] loss: 7.220839, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.614336
[epoch: 1, batch: 44] loss: 6.997747, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.391244
[epoch: 1, batch: 45] loss: 7.289024, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.682521
[epoch: 1, batch: 46] loss: 8.104959, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.498456
[epoch: 1, batch: 47] loss: 7.818621, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.212118
[epoch: 1, batch: 48] loss: 7.451015, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.844513
[epoch: 1, batch: 49] loss: 7.190177, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.583674
[epoch: 1, batch: 50] loss: 7.414326, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.807823
[epoch: 1, batch: 51] loss: 7.082467, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.475964
[epoch: 1, batch: 52] loss: 8.054209, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.447706
[epoch: 1, batch: 53] loss: 7.361894, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.755390
[epoch: 1, batch: 54] loss: 7.404776, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.798274
[epoch: 1, batch: 55] loss: 7.673492, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.066990
[epoch: 1, batch: 56] loss: 7.255838, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.649335
[epoch: 1, batch: 57] loss: 7.364078, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.757576
[epoch: 1, batch: 58] loss: 7.672638, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.066136
[epoch: 1, batch: 59] loss: 6.943457, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.336953
[epoch: 1, batch: 60] loss: 7.074402, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.467899
[epoch: 1, batch: 61] loss: 7.573625, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.967122
[epoch: 1, batch: 62] loss: 6.797518, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.191013
[epoch: 1, batch: 63] loss: 7.667085, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  2.060580
[epoch: 1, batch: 64] loss: 7.461186, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.854683
[epoch: 1, batch: 65] loss: 7.218839, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.612336
[epoch: 1, batch: 66] loss: 6.966671, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.360166
[epoch: 1, batch: 67] loss: 7.419224, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.812721
[epoch: 1, batch: 68] loss: 7.038870, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.432365
[epoch: 1, batch: 69] loss: 7.527884, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.921382
[epoch: 1, batch: 70] loss: 6.735312, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.128809
[epoch: 1, batch: 71] loss: 7.051725, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.445222
[epoch: 1, batch: 72] loss: 7.819386, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.212884
[epoch: 1, batch: 73] loss: 7.343047, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.736544
[epoch: 1, batch: 74] loss: 6.986712, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.380208
[epoch: 1, batch: 75] loss: 7.757770, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.151267
[epoch: 1, batch: 76] loss: 7.547459, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.940956
[epoch: 1, batch: 77] loss: 7.939866, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.333363
[epoch: 1, batch: 78] loss: 7.063341, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.456838
[epoch: 1, batch: 79] loss: 7.327311, obj_loss:  1.039721, cls_loss:  4.566783, reg_loss:  1.720806
[epoch: 1, batch: 80] loss: 7.129947, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.523444
[epoch: 1, batch: 81] loss: 7.087235, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.480733
[epoch: 1, batch: 82] loss: 7.572030, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.965527
[epoch: 1, batch: 83] loss: 7.192766, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.586263
[epoch: 1, batch: 84] loss: 7.427009, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.820506
[epoch: 1, batch: 85] loss: 7.694969, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.088466
[epoch: 1, batch: 86] loss: 8.016901, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  2.410399
[epoch: 1, batch: 87] loss: 7.126433, obj_loss:  1.039721, cls_loss:  4.566782, reg_loss:  1.519930
</code></pre>
<p>It’s the first epoch. Loss obj_loss and cls_loss is unlikely to be optimized after batch 26, yet reg_loss remains shifting.<br>
The training script:</p>
<pre><code class="lang-auto">    device = torch.device("cuda")
    dataloader_train = get_dataloader(train=True, batch_size=batch_size)
    dataset_val = VOCDetectionDataset(root="/data/sfy_projects/Datasets/VOC2007/VOCtrainval_06-Nov-2007", year="2007",
                                      image_set="val",
                                      transforms=get_transforms(train=False),
                                      show=False)
    coco_anno_path = "/data/sfy_projects/Datasets/VOC2007/VOCtrainval_06-Nov-2007/voc2007.val.cocoformat.json"

    model = YOLO(anchors=k_means_anchors["2007.train"], num_classes=20).to(device)
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-5)
    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=5, verbose=True)

    for epoch in range(epochs):
        print("----------------------  TRAINING  ---------------------- ")
        model.train()
        model.freeze_backbone(layers=3)
        model.freeze_bn(model.backbone)
        running_loss, running_objectness_loss, running_classification_loss, running_regression_loss = 0.0, 0.0, 0.0, 0.0
        for i, data in enumerate(dataloader_train):
            imgs, targets = data
            imgs = imgs.to(device)
            for target in targets:
                for obj in target["objects"]:
                    obj["class"] = obj["class"].to(device)
                    obj["bbox"] = obj["bbox"].to(device)

            res = model(imgs, targets)
            loss = res["loss"]["loss"]
            objectness_loss = res["loss"]["objectness_loss"]
            classification_loss = res["loss"]["classification_loss"]
            regression_loss = res["loss"]["regression_loss"]
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            running_objectness_loss += objectness_loss.item()
            running_classification_loss += classification_loss.item()
            running_regression_loss += regression_loss.item()
            if (i + 1) % batches_show == 0:
                print(f"[epoch: {epoch + 1}, batch: {i + 1}] loss: {running_loss / batches_show:.6f}, obj_loss: {running_objectness_loss / batches_show: .6f}, cls_loss: {running_classification_loss / batches_show: .6f}, reg_loss: {running_regression_loss / batches_show: .6f}")
                running_loss, running_objectness_loss, running_classification_loss, running_regression_loss = 0.0, 0.0, 0.0, 0.0
        lr_scheduler.step(loss)
</code></pre>
<p>I’ve tried to ommit optimzing regression loss. Finally I got a rather good model which points out the object’s location and class without awareness of it’s bbox size. But… uhm, I’m not sure what can this prove.<br>
My biggest doubt is that, since I already called <code>optimizer.zero_grad()</code>, why can something happend in one batch of training have a lasting effect on the following training batches? Also I cannot find any clue why this happens.<br>
BTW, the problem is reproducible to me. In fact, it happens at the first epoch every time if I set the learning rate as 0.01. However the sudden raise of regression loss never show up again yet since I tuned learning rate to 0.001(as far as I’ve seen). Is all of this just about learning rate?</p>