-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathtrainer.py
134 lines (109 loc) · 4.44 KB
/
trainer.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
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
import torch
import copy
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.autograd import Variable
from tqdm import tqdm
from datetime import datetime, timezone
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
entity="adorable-lantanas"
project="bp-vs-wp"
prefix="tbishnoi-run-0"
log_results=True
num_ramdom_seed=10
random_seed=0 # initial random_seed
_configs = dict(
entity=entity,
project=project,
epochs = 10,
batch_size = 32,
num_inputs = 784,
num_hidden = 100,
num_outputs = 10,
activation_type = 'relu',
bias=True,
momentum=0.9,
weight_decay=0.000,
nesterov=True,
)
"""
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
/$$ /$$ /$$ /$$
| $$ |__/ |__/ | $$
/$$$$$$ /$$$$$$ /$$$$$$ /$$ /$$$$$$$ /$$ /$$$$$$$ /$$$$$$ /$$$$$$ | $$ /$$$$$$ /$$$$$$$
|_ $$_/ /$$__ $$|____ $$| $$| $$__ $$| $$| $$__ $$ /$$__ $$ /$$__ $$| $$ |____ $$| $$__ $$
| $$ | $$ \__/ /$$$$$$$| $$| $$ \ $$| $$| $$ \ $$| $$ \ $$ | $$ \ $$| $$ /$$$$$$$| $$ \ $$
| $$ /$$| $$ /$$__ $$| $$| $$ | $$| $$| $$ | $$| $$ | $$ | $$ | $$| $$ /$$__ $$| $$ | $$
| $$$$/| $$ | $$$$$$$| $$| $$ | $$| $$| $$ | $$| $$$$$$$ | $$$$$$$/| $$| $$$$$$$| $$ | $$
\___/ |__/ \_______/|__/|__/ |__/|__/|__/ |__/ \____ $$ | $$____/ |__/ \_______/|__/ |__/
/$$ \ $$ | $$
| $$$$$$/ | $$
\______/ |__/
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
"""
def main():
global random_seed
from utils.data_utils import set_seed, download_mnist
# set initial random_seed
set_seed(random_seed)
#create loaders
train_set, valid_set, test_set = download_mnist()
configs = copy.deepcopy(_configs)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=configs['batch_size'], shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_set, batch_size=configs['batch_size'], shuffle=False)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=configs['batch_size'], shuffle=False)
from utils.training_utils import select_model, train_model
from rules.classes.BasicOptim import BasicOptimizer
for rule in [
'backprop',
# 'hebb',
'wp',
# 'np',
]:
# for each rule we copy base config
configs = copy.deepcopy(_configs)
# conditionally modify configs
if rule=='backprop':
configs['lr'] = 1e-2
elif rule=='wp':
configs['lr'] = 1e-4
configs['sigma'] = 1e-4
else:
raise NotImplementedError(f"rule: {rule} outside implementation of training plan!")
for random_seed in tqdm(range(num_ramdom_seed)):
# rule
configs['rule_select'] = rule
# experiment name
experiment_name = f"{prefix}-{rule}-{random_seed}"
configs['experiment_name'] = experiment_name
# create and save model name
model_filepath = f"models/model-{datetime.now(timezone.utc).strftime('%y%m%d-%H%M%S')}.pth"
configs['model_filepath'] = model_filepath
# overwrites the random_seed that was used before
set_seed(random_seed)
configs['random_seed'] = random_seed
# select model
model = select_model(configs, device)
optimizer = BasicOptimizer(model.parameters(), lr=configs['lr'], weight_decay=configs['weight_decay'])
print()
print(f"\texp name:\t{experiment_name}")
print(f"\ttrain rule:\t{rule}")
print(f"\tnum epochs:\t{configs['epochs']}")
print(f"\trandom seed:\t{random_seed}")
print(f"\tmodel filepath:\t{model_filepath}")
print()
train_model(
model,
train_loader,
test_loader,
optimizer,
experiment_name,
configs,
log_results=log_results,
device=device
)
del model
del optimizer
if __name__ == '__main__':
main()