GNN Irregular loss values (distorted)

I’m working with different datasets from MoleculeNet but for all I’m getting irregularity in loss, the loss is increasing and decreasing instantly.

I tried different layers, learning rates, and even tried different datasets from MoleculeNet but, I’m getting almost the same result, this distortion increases but never decreases.

How can I improve this?

output

Training Started...!
Epoch 0 | Train Loss 1.3184049129486084
Epoch 10 | Train Loss 0.8975121974945068
Epoch 20 | Train Loss 4.639061450958252
Epoch 30 | Train Loss 1.0549440383911133
Epoch 40 | Train Loss 2.291563034057617
Epoch 50 | Train Loss 0.40543055534362793
Epoch 60 | Train Loss 3.7577757835388184
Epoch 70 | Train Loss 2.152988910675049
Epoch 80 | Train Loss 0.993300199508667
Epoch 90 | Train Loss 2.9262337684631348
Epoch 100 | Train Loss 0.40380847454071045
Epoch 110 | Train Loss 1.1322712898254395
Epoch 120 | Train Loss 1.6807966232299805
Epoch 130 | Train Loss 2.4824788570404053
Epoch 140 | Train Loss 2.337397575378418
Epoch 150 | Train Loss 2.8971199989318848
Epoch 160 | Train Loss 2.8942408561706543
Epoch 170 | Train Loss 1.8320939540863037
Epoch 180 | Train Loss 1.6915087699890137
Epoch 190 | Train Loss 0.29916471242904663
Epoch 200 | Train Loss 1.306685209274292
Epoch 210 | Train Loss 3.2683615684509277
Epoch 220 | Train Loss 0.06050372123718262
Epoch 230 | Train Loss 0.38682126998901367
Epoch 240 | Train Loss 1.1061413288116455
Epoch 250 | Train Loss 0.22812628746032715
Epoch 260 | Train Loss 0.6184450387954712
Epoch 270 | Train Loss 1.4897408485412598
Epoch 280 | Train Loss 4.929813861846924
Epoch 290 | Train Loss 0.11922299861907959
Epoch 300 | Train Loss 3.699568271636963
Epoch 310 | Train Loss 0.02466607093811035
Epoch 320 | Train Loss 1.8530726432800293
Epoch 330 | Train Loss 0.8568987846374512
Epoch 340 | Train Loss 0.36516761779785156
Epoch 350 | Train Loss 0.9791746139526367
Epoch 360 | Train Loss 0.37069225311279297
Epoch 370 | Train Loss 1.3499877452850342
Epoch 380 | Train Loss 0.519951343536377
Epoch 390 | Train Loss 3.8632442951202393
Epoch 400 | Train Loss 0.4643235206604004
Epoch 410 | Train Loss 0.8706831932067871
Epoch 420 | Train Loss 1.1726784706115723
Epoch 430 | Train Loss 1.2985022068023682
Epoch 440 | Train Loss 7.224368095397949
Epoch 450 | Train Loss 1.1312479972839355
Epoch 460 | Train Loss 0.3077009916305542
Epoch 470 | Train Loss 0.6891758441925049
Epoch 480 | Train Loss 1.6510722637176514
Epoch 490 | Train Loss 0.7090321779251099
Epoch 500 | Train Loss 0.8605406284332275
Epoch 510 | Train Loss 1.3280174732208252
Epoch 520 | Train Loss 0.9676799774169922
Epoch 530 | Train Loss 1.5659568309783936
Epoch 540 | Train Loss 0.7733702659606934
Epoch 550 | Train Loss 0.7845138907432556
Epoch 560 | Train Loss 1.72637939453125
Epoch 570 | Train Loss 0.7738587856292725
Epoch 580 | Train Loss 1.0326557159423828
Epoch 590 | Train Loss 5.928122520446777
Epoch 600 | Train Loss 2.1471190452575684
Epoch 610 | Train Loss 0.5509365200996399
Epoch 620 | Train Loss 0.7522091865539551
Epoch 630 | Train Loss 1.0490703582763672
Epoch 640 | Train Loss 3.6967601776123047
Epoch 650 | Train Loss 0.04169821739196777
Epoch 660 | Train Loss 0.16604477167129517
Epoch 670 | Train Loss 0.34248971939086914
Epoch 680 | Train Loss 0.21051549911499023
Epoch 690 | Train Loss 1.3220789432525635
Epoch 700 | Train Loss 1.8236854076385498
Epoch 710 | Train Loss 4.991037368774414
Epoch 720 | Train Loss 0.09024512767791748
Epoch 730 | Train Loss 0.7814116477966309
Epoch 740 | Train Loss 0.23316997289657593
Epoch 750 | Train Loss 0.7268493175506592
Epoch 760 | Train Loss 0.4061264991760254
Epoch 770 | Train Loss 1.0022287368774414
Epoch 780 | Train Loss 1.9513254165649414
Epoch 790 | Train Loss 0.6404123306274414
Epoch 800 | Train Loss 0.8553285598754883
Epoch 810 | Train Loss 0.3502233028411865
Epoch 820 | Train Loss 1.8192477226257324
Epoch 830 | Train Loss 3.4299943447113037
Epoch 840 | Train Loss 0.43554258346557617
Epoch 850 | Train Loss 0.7778606414794922
Epoch 860 | Train Loss 2.384077548980713
Epoch 870 | Train Loss 0.5119071006774902
Epoch 880 | Train Loss 0.5938878059387207
Epoch 890 | Train Loss 0.178863525390625
Epoch 900 | Train Loss 0.9071254730224609
Epoch 910 | Train Loss 2.315983772277832
Epoch 920 | Train Loss 0.3971567451953888
Epoch 930 | Train Loss 0.2216503620147705
Epoch 940 | Train Loss 1.4241361618041992
Epoch 950 | Train Loss 0.9637584686279297
Epoch 960 | Train Loss 1.0516422986984253
Epoch 970 | Train Loss 0.7053394317626953
Epoch 980 | Train Loss 0.44690942764282227
Epoch 990 | Train Loss 0.7075588703155518
Epoch 1000 | Train Loss 1.5911892652511597
Epoch 1010 | Train Loss 1.4753718376159668
Epoch 1020 | Train Loss 1.3572875261306763
Epoch 1030 | Train Loss 0.670606791973114
Epoch 1040 | Train Loss 0.7239326238632202
Epoch 1050 | Train Loss 1.328798770904541
Epoch 1060 | Train Loss 0.24600374698638916
Epoch 1070 | Train Loss 0.746793270111084
Epoch 1080 | Train Loss 0.7242438793182373
Epoch 1090 | Train Loss 0.8730747699737549
Epoch 1100 | Train Loss 0.3874173164367676
Epoch 1110 | Train Loss 1.380805492401123
Epoch 1120 | Train Loss 0.062389373779296875
Epoch 1130 | Train Loss 3.2198665142059326
Epoch 1140 | Train Loss 0.2733985483646393
Epoch 1150 | Train Loss 0.4619467258453369
Epoch 1160 | Train Loss 0.3873324394226074
Epoch 1170 | Train Loss 1.7770310640335083
Epoch 1180 | Train Loss 0.857968807220459
Epoch 1190 | Train Loss 1.8375463485717773
Epoch 1200 | Train Loss 1.042975902557373
Epoch 1210 | Train Loss 0.09935760498046875
Epoch 1220 | Train Loss 2.288851737976074
Epoch 1230 | Train Loss 1.470461368560791
Epoch 1240 | Train Loss 1.1549944877624512
Epoch 1250 | Train Loss 0.37267041206359863
Epoch 1260 | Train Loss 0.6693940162658691
Epoch 1270 | Train Loss 0.2859673500061035
Epoch 1280 | Train Loss 0.8468062877655029
Epoch 1290 | Train Loss 0.05202460289001465
Epoch 1300 | Train Loss 0.3275785446166992
Epoch 1310 | Train Loss 0.9826569557189941
Epoch 1320 | Train Loss 0.47000551223754883
Epoch 1330 | Train Loss 0.8779596090316772
Epoch 1340 | Train Loss 0.4970674514770508
Epoch 1350 | Train Loss 0.07213175296783447
Epoch 1360 | Train Loss 0.0683889389038086
Epoch 1370 | Train Loss 1.538172721862793
Epoch 1380 | Train Loss 0.17601919174194336
Epoch 1390 | Train Loss 0.4858182966709137
Epoch 1400 | Train Loss 0.06339013576507568
Epoch 1410 | Train Loss 1.210093379020691
Epoch 1420 | Train Loss 0.4947328567504883
Epoch 1430 | Train Loss 0.42642533779144287
Epoch 1440 | Train Loss 0.1839771270751953
Epoch 1450 | Train Loss 0.2670316696166992
Epoch 1460 | Train Loss 0.1603527069091797
Epoch 1470 | Train Loss 1.3721613883972168
Epoch 1480 | Train Loss 0.8607507348060608
Epoch 1490 | Train Loss 0.7421789169311523
Epoch 1500 | Train Loss 0.08590173721313477
Epoch 1510 | Train Loss 0.12283754348754883
Epoch 1520 | Train Loss 0.21935272216796875
Epoch 1530 | Train Loss 1.1690268516540527
Epoch 1540 | Train Loss 0.6230419874191284
Epoch 1550 | Train Loss 0.10403728485107422
Epoch 1560 | Train Loss 0.20833683013916016
Epoch 1570 | Train Loss 0.17981243133544922
Epoch 1580 | Train Loss 1.0952425003051758
Epoch 1590 | Train Loss 0.8331184387207031
Epoch 1600 | Train Loss 1.2511465549468994
Epoch 1610 | Train Loss 3.0601024627685547
Epoch 1620 | Train Loss 0.14174747467041016
Epoch 1630 | Train Loss 0.08653402328491211
Epoch 1640 | Train Loss 0.07086014747619629
Epoch 1650 | Train Loss 0.9229309558868408
Epoch 1660 | Train Loss 1.0049333572387695
Epoch 1670 | Train Loss 1.9561171531677246
Epoch 1680 | Train Loss 0.0049686431884765625
Epoch 1690 | Train Loss 0.6676177978515625
Epoch 1700 | Train Loss 0.7200961112976074
Epoch 1710 | Train Loss 0.44962120056152344
Epoch 1720 | Train Loss 0.19601678848266602
Epoch 1730 | Train Loss 0.8548879623413086
Epoch 1740 | Train Loss 1.7700645923614502
Epoch 1750 | Train Loss 0.23295307159423828
Epoch 1760 | Train Loss 0.09788763523101807
Epoch 1770 | Train Loss 3.442563533782959
Epoch 1780 | Train Loss 0.3740963935852051
Epoch 1790 | Train Loss 0.14022612571716309
Epoch 1800 | Train Loss 1.0260200500488281
Epoch 1810 | Train Loss 0.32472580671310425
Epoch 1820 | Train Loss 0.5827023983001709
Epoch 1830 | Train Loss 0.1571100950241089
Epoch 1840 | Train Loss 0.9010820388793945
Epoch 1850 | Train Loss 0.37263745069503784
Epoch 1860 | Train Loss 0.9629731178283691
Epoch 1870 | Train Loss 0.5000338554382324
Epoch 1880 | Train Loss 0.19631600379943848
Epoch 1890 | Train Loss 0.5523502826690674
Epoch 1900 | Train Loss 0.09037113189697266
Epoch 1910 | Train Loss 1.1198182106018066
Epoch 1920 | Train Loss 0.6817159652709961
Epoch 1930 | Train Loss 0.8300008773803711
Epoch 1940 | Train Loss 0.39383769035339355
Epoch 1950 | Train Loss 0.7858984470367432
Epoch 1960 | Train Loss 0.0002378523349761963
Epoch 1970 | Train Loss 0.12754535675048828
Epoch 1980 | Train Loss 0.9645538330078125
Epoch 1990 | Train Loss 0.045568645000457764
Training Completed...!
Model Saved...!

Here’s the model:

n_nn = 64

class Net(nn.Module):
    def __init__(self, n_feats) -> None:
        super().__init__()
        T.manual_seed(42)

        self.input = gnn.GCNConv(n_feats, n_nn)
        self.conv1 = gnn.GCNConv(n_nn, n_nn)
        self.conv2 = gnn.GCNConv(n_nn, n_nn)
        self.conv3 = gnn.GCNConv(n_nn, n_nn)
        self.out = nn.Linear(n_nn * 2, 1)

    def forward(self, batch_data):
        x, edge_index, batch = batch_data.x.float(), batch_data.edge_index, batch_data.batch

        x = T.tanh(self.input(x, edge_index))
        x = T.tanh(self.conv1(x, edge_index))
        x = T.tanh(self.conv2(x, edge_index))
        x = T.tanh(self.conv3(x, edge_index))

        return self.out(
            T.cat([gnn.global_mean_pool(x, batch), gnn.global_max_pool(x, batch)], dim=1))

Training Model:

class Train:

    def __init__(self, model, train_batch) -> None:
        self.model = model
        self.train_batch = train_batch
        self.loss_fn = nn.MSELoss()
        self.optimizer = optim.Adam(model.parameters(), lr=0.0007)
        self.losses = []

    def start_training(self, epochs=2000):
        self.losses.clear()
        print("Training Started...!")
        for epoch in range(epochs):
            loss = self._train()
            self.losses.append(loss)
            if epoch % 10 == 0:
                print(f"Epoch {epoch} | Train Loss {loss}")
        print("Training Completed...!")
        T.save(self.model.state_dict(), model_path)
        print("Model Saved...!")

    def _train(self):
        loss = 0
        for batch in self.train_batch:
            self.optimizer.zero_grad()
            out = self.model(batch)
            loss = T.sqrt(self.loss_fn(out, batch.y))
            loss.backward()
            self.optimizer.step()
        return loss```