-
Notifications
You must be signed in to change notification settings - Fork 0
/
exp_SNN.py
69 lines (55 loc) · 2.21 KB
/
exp_SNN.py
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
#!/usr/bin/env python
#SBATCH --job-name=SNN
#SBATCH --error=%x.%j.err
#SBATCH --output=%x.%j.out
#SBATCH [email protected]
#SBATCH --export=ALL
#SBATCH --time=48:00:00
#SBATCH --partition=sdil
#SBATCH --gres=gpu:1
import os
import sys
sys.path.append(os.getcwd())
sys.path.append(os.path.join(os.getcwd(), 'utils'))
from configuration import *
import torch
import pprint
import BaselineSNN as B
from utils import *
args = parser.parse_args()
for seed in range(10):
args.SEED = seed
args = FormulateArgs(args)
print(f'Training network on device: {args.DEVICE}.')
MakeFolder(args)
train_loader, datainfo = GetDataLoader(args, 'train')
valid_loader, datainfo = GetDataLoader(args, 'valid')
test_loader, datainfo = GetDataLoader(args, 'test')
pprint.pprint(datainfo)
SetSeed(args.SEED)
setup = f"model_SNN_data_{datainfo['dataname']}_seed_{args.SEED:02d}.model"
print(f'Training setup: {setup}.')
msglogger = GetMessageLogger(args, setup)
msglogger.info(f'Training network on device: {args.DEVICE}.')
msglogger.info(f'Training setup: {setup}.')
msglogger.info(datainfo)
if os.path.isfile(f'{args.savepath}/{setup}'):
print(f'{setup} exists, skip this training.')
msglogger.info('Training was already finished.')
else:
topology = [datainfo['N_feature']] + args.hidden + [datainfo['N_class']]
msglogger.info(f'Topology of the network: {topology}.')
snn = B.SpikingNeuralNetwork(args, topology).to(args.DEVICE)
lossfunction = B.SNNLoss().to(args.DEVICE)
optimizer = torch.optim.Adam(snn.GetParam(), lr=args.LR)
if args.PROGRESSIVE:
snn, best = train_pnn_progressive(snn, train_loader, valid_loader, lossfunction, optimizer, args, msglogger, UUID=setup)
else:
snn, best = train_pnn(snn, train_loader, valid_loader, lossfunction, optimizer, args, msglogger, UUID=setup)
if best:
if not os.path.exists(f'{args.savepath}/'):
os.makedirs(f'{args.savepath}/')
torch.save(snn, f'{args.savepath}/{setup}')
msglogger.info('Training if finished.')
else:
msglogger.warning('Time out, further training is necessary.')