-
Notifications
You must be signed in to change notification settings - Fork 43
/
train.py
59 lines (47 loc) · 1.86 KB
/
train.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
'''
@Author: Zhou Kai
@GitHub: https://github.com/athon2
@Date: 2018-11-30 09:53:44
'''
import torch
from torch.autograd import Variable
from tqdm import tqdm
from utils import AverageMeter, calculate_accuracy
def train_epoch(epoch, data_set, model, criterion, optimizer, opt, logger):
print('train at epoch {}'.format(epoch))
model.train()
losses = AverageMeter()
accuracies = AverageMeter()
data_set.file_open()
train_loader = torch.utils.data.DataLoader(dataset=data_set,
batch_size=opt["batch_size"],
shuffle=True,
pin_memory=True)
training_process = tqdm(train_loader)
for i, (inputs, targets) in enumerate(training_process):
if i > 0:
training_process.set_description("Loss: %.4f, Acc: %.4f"%(losses.avg.item(), accuracies.avg.item()))
if opt["cuda_devices"] is not None:
inputs = inputs.type(torch.FloatTensor)
inputs = inputs.cuda()
targets = targets.type(torch.FloatTensor)
targets = targets.cuda()
if opt["VAE_enable"]:
outputs, distr = model(inputs)
loss = criterion(outputs, targets, distr)
else:
outputs = model(inputs)
loss = criterion(outputs, targets)
acc = calculate_accuracy(outputs.cpu(), targets.cpu())
losses.update(loss.cpu(), inputs.size(0))
accuracies.update(acc, inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
logger.log(phase="train",values={
'epoch': epoch,
'loss': format(losses.avg.item(), '.4f'),
'acc': format(accuracies.avg.item(), '.4f'),
'lr': optimizer.param_groups[0]['lr']
})
data_set.file_close()