SkipResConnection/train_conv_snndirect.py at main · Intelligent-Computing-Lab-Panda/SkipResConnection · 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
import config
from models.model_snndirect import *
# from model import *


from utils import *

torch.manual_seed(1234)
torch.cuda.manual_seed(1234)
device = torch.device("cuda")
cudnn.benchmark = True
cudnn.deterministic = True


def main(args):

    print (args.t_init)
    args = config.get_args()

    train_loader, test_loader = data_load(args)

    criterion = torch.nn.CrossEntropyLoss()
    weight_sum_K = 1
    scale = args.t_scale

    print ("t scale:", scale)

    if args.arch == 'base':
        model = mid_vgg_direct(max_t=scale)
    elif args.arch == 'res':
        model = mid_vgg_direct_residual(max_t=scale)
    elif args.arch == 'shuffle':
        model = mid_shufflenet_direct(max_t=scale, t_init = args.t_init)

    model = model.cuda()


    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,weight_decay=1e-3)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs,eta_min=0)


    losses = AverageMeter('Loss', ':.4f')
    progress = ProgressMeter(
        len(train_loader),
        [losses])


    all_testacclist = []
    all_trainacclist = []
    all_losslist = []
    all_delaymeanlist = []
    all_delaystdlist = []

    delay_mean = torch.mean(torch.stack([model.t_shift1.flatten(), model.t_shift2.flatten()]))
    delay_std = torch.std(torch.stack([model.t_shift1.flatten(), model.t_shift2.flatten()]))
    print ('init delay..')
    print (delay_mean, delay_std)
    for epoch in range(args.epochs):
        model.train()

        epoch_loss = []
        for i, (data, target) in enumerate(train_loader):
            data = data.cuda()
            target = target.cuda()
            optimizer.zero_grad()

            if args.arch == 'shuffle':
                outs,weight_sum_cost, force_loss = model(data)
            else:
                outs,weight_sum_cost  = model(data)
                force_loss = 0

            ce_loss = criterion(-1*(outs), target)

            reg_loss = weight_sum_K*weight_sum_cost
            loss = ce_loss+ reg_loss + force_loss*1e-6
            losses.update((torch.mean(ce_loss)).item(), data.shape[0])

            loss.sum().backward()
            epoch_loss.append(float(ce_loss.data.cpu().numpy()))

            optimizer.step()

            if (i+1) % 1500 == 0:    # print every 2000 mini-batches
                progress.display(i)
        scheduler.step()

        all_losslist.append(sum(epoch_loss)/len(epoch_loss))
        if (epoch+1) % 5 ==0:
            acc = test(model, test_loader, epoch,args)
            acc_train = test(model, train_loader, epoch,args)

            all_testacclist.append(acc)
            all_trainacclist.append(acc_train)

            if args.arch == 'shuffle':
                delay_mean = torch.mean(torch.stack([model.t_shift1.flatten(), model.t_shift2.flatten()]))
                delay_std = torch.std(torch.stack([model.t_shift1.flatten(), model.t_shift2.flatten()]))
                all_delaymeanlist.append(float(delay_mean.cpu().data.numpy()))
                all_delaystdlist.append(float(delay_std.cpu().data.numpy()))

    print (all_losslist)
    print(all_testacclist)

    torch.save(model.state_dict(), 'savemodel/{}_{}_init{}_t{}_b{}_lr{}_ep{}'.format(args.dataset, args.arch, args.t_init, args.t_scale, args.batch_size, args.lr, args.epochs))

    return np.max(all_testacclist)

def test(model, test_loader, epoch,args):

    model.eval()
    correct = 0
    total = 0
    first_spike_time_list = []

    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)

            # the class with the highest energy is what we choose as prediction
            first_spike_time, predicted = torch.min(outputs.data, 1)
            first_spike_time_list.append(torch.log(first_spike_time).sum() / first_spike_time.size(0))
            total += labels.size(0)
            correct += (predicted == labels).sum().item()


    print("first_spike mean time:", sum(first_spike_time_list) / len(first_spike_time_list))


    return (100 * correct / total)

if __name__ == '__main__':
    args = config.get_args()
    main(args)