SC-Wireless/main_train.py at main · Restuccia-Group/SC-Wireless · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import numpy as np
import os
from matplotlib import pyplot as plt

import torch
import torch.nn as nn

import time

import arg_parser
from datasets import setup_dataloaders
from utils import *

from trainer import *
from losses import *
from utils import load_state_dict


best_sa = 0


def main():
    global args, best_sa
    args = arg_parser.parse_args()

    torch.cuda.set_device(int(args.gpu))
    os.makedirs(args.save_dir, exist_ok=True)
    if args.seed:
        setup_seed(args.seed)

    # prepare dataset
    (
        model,
        train_loader,
        val_loader,
        test_loader,
    ) = setup_model_dataset(args)
    model.cuda()

    print(f"number of train dataset {len(train_loader.dataset)}")
    print(f"number of val dataset {len(val_loader.dataset)}")


    if args.loss=="entropy_regularizer":
        criterion = entropy_loss
    elif args.loss=="supcon":
        criterion = supcon_loss
    else:
        criterion = nn.CrossEntropyLoss()

    decreasing_lr = list(map(int, args.decreasing_lr.split(",")))


    optimizer, scheduler = get_optimizer_and_scheduler(model, args)

    all_result = {}
    all_result["train_ta"] = []
    all_result["test_ta"] = []
    all_result["val_ta"] = []

    start_epoch = 0
    state = 0

    swa_model = torch.optim.swa_utils.AveragedModel(model)
    swa_start = -1
    swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, swa_lr=0.0001)


    for epoch in range(start_epoch, args.epochs):
        start_time = time.time()
        print(
            "Epoch #{}, Learning rate: {}".format(
                epoch, optimizer.state_dict()["param_groups"][0]["lr"]
            )
        )
        acc = train(train_loader, model, criterion, optimizer, epoch, args)

        if epoch > swa_start and swa_start!=-1:
            swa_model.update_parameters(model)
            swa_scheduler.step()
        else:
            scheduler.step()

        # evaluate on validation set
        tacc = validate(val_loader, model, criterion, args)
        # scheduler.step()

        all_result["train_ta"].append(acc)
        all_result["val_ta"].append(tacc)

        # remember best prec@1 and save checkpoint
        is_best_sa = tacc > best_sa
        best_sa = max(tacc, best_sa)

        save_checkpoint(
            {
                "result": all_result,
                "epoch": epoch + 1,
                "state_dict": model.state_dict(),
                "best_sa": best_sa,
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
            },
            is_SA_best=is_best_sa,
            save_path=args.save_dir,
            args=args
        )
        print("one epoch duration:{}".format(time.time() - start_time))

    if swa_start>-1:
        # swa_model = load_state_dict(swa_model, args, filename="model_SA_best.pth.tar")
        torch.optim.swa_utils.update_bn(train_loader, swa_model)

        save_checkpoint(
                {
                    "result": all_result,
                    "epoch": epoch + 1,
                    "state_dict": swa_model.state_dict(),
                    "best_sa": best_sa,
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                },
                is_SA_best=is_best_sa,
                save_path=args.save_dir,
                args=args
            )


    # plot training curve
    plt.plot(all_result["train_ta"], label="train_acc")
    plt.plot(all_result["val_ta"], label="val_acc")
    plt.legend()
    plt.savefig(os.path.join(args.save_dir, str(state) + "net_train.png"))
    plt.close()

    print("Performance on the test data set")

    try:
        test_acc = validate(test_loader, model, criterion, args)
    except:
        test_acc = validate(val_loader, model, criterion, args)

    if len(all_result["val_ta"]) != 0:
        val_pick_best_epoch = np.argmax(np.array(all_result["val_ta"]))
        print(
            "* best SA = {}, Epoch = {}".format(
                all_result["val_ta"][val_pick_best_epoch], val_pick_best_epoch + 1
            )
        )

if __name__ == "__main__":
    main()