-
Notifications
You must be signed in to change notification settings - Fork 2
/
train_utils.py
140 lines (110 loc) · 4.71 KB
/
train_utils.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
135
136
137
138
139
import os
import shutil
import time
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
from torch.optim import SGD, Optimizer
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from architectures import get_architecture
from datasets import get_dataset, get_num_classes
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified values of k"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def init_logfile(filename: str, text: str):
f = open(filename, 'w')
f.write(text+"\n")
f.close()
def log(filename: str, text: str):
f = open(filename, 'a')
f.write(text+"\n")
f.close()
def requires_grad_(model:torch.nn.Module, requires_grad:bool) -> None:
for param in model.parameters():
param.requires_grad_(requires_grad)
def copy_code(outdir):
"""Copies files to the outdir to store complete script with each experiment"""
# embed()
code = []
exclude = set([])
for root, _, files in os.walk("./code", topdown=True):
for f in files:
if not f.endswith('.py'):
continue
code += [(root,f)]
for r, f in code:
codedir = os.path.join(outdir,r)
if not os.path.exists(codedir):
os.mkdir(codedir)
shutil.copy2(os.path.join(r,f), os.path.join(codedir,f))
print("Code copied to '{}'".format(outdir))
def prologue(args):
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
# Copies files to the outdir to store complete script with each experiment
copy_code(args.outdir)
train_dataset = get_dataset(args.dataset, 'train')
test_dataset = get_dataset(args.dataset, 'test')
pin_memory = (args.dataset == "imagenet")
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch,
num_workers=args.workers, pin_memory=pin_memory)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=args.batch,
num_workers=args.workers, pin_memory=pin_memory)
if args.pretrained_model != '':
# assert args.arch == 'cifar_resnet110', 'Unsupported architecture for pretraining'
checkpoint = torch.load(args.pretrained_model)
model = get_architecture(checkpoint["arch"], args.dataset)
model.load_state_dict(checkpoint['state_dict'])
# model[1].fc = nn.Linear(64, get_num_classes('cifar10')).to(device)
else:
model = get_architecture(args.arch, args.dataset)
logfilename = os.path.join(args.outdir, 'log.txt')
init_logfile(logfilename, "epoch\ttime\tlr\ttrain loss\ttrain acc\ttestloss\ttest acc")
writer = SummaryWriter(args.outdir)
criterion = CrossEntropyLoss().to(device)
optimizer = SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = StepLR(optimizer, step_size=args.lr_step_size, gamma=args.gamma)
starting_epoch = 0
# Load latest checkpoint if exists (to handle philly failures)
model_path = os.path.join(args.outdir, 'checkpoint.pth.tar')
if args.resume:
if os.path.isfile(model_path):
print("=> loading checkpoint '{}'".format(model_path))
checkpoint = torch.load(model_path,
map_location=lambda storage, loc: storage)
starting_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(model_path, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(model_path))
return train_loader, test_loader, criterion, model, optimizer, scheduler, \
starting_epoch, logfilename, model_path, device, writer