forked from mingliangzhang2018/AliProducts-Challenge
-
Notifications
You must be signed in to change notification settings - Fork 0
/
function.py
94 lines (77 loc) · 3.2 KB
/
function.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
import numpy as np
import torch
import time
from torch.autograd import Variable
def compute_top5(target, output):
y_resize = target.view(-1, 1)
correct_top5 = np.zeros(5)
for maxk in range(1, 6):
_, pred = output.topk(maxk, 1)
correct_top5[maxk - 1] = torch.eq(pred, y_resize).sum().float().item()
return correct_top5
def train_model(trainLoader, model, epoch, optimizer, criterion, cfg):
model.train()
# 损失及正确率累加变量初始化
running_loss = 0
running_correct = 0
start_time = time.time()
for batch_idx, (data, target) in enumerate(trainLoader):
# 针对Cuda进行设置
data, target = data.cuda(), target.cuda()
# 训练阶段优化器梯度初始化
optimizer.zero_grad()
# 模型前向传播
output = model(data)
# 得到训练误差
loss = criterion(output, target)
# 记录误差及准确率
running_loss += loss.data
# running_correct_top5 += compute_top5(target, output)
pred = output.data.max(1, keepdim=True)[1] # 获得得分最高的类别
running_correct += pred.eq(target.data.view_as(pred)).cpu().sum()
# 在训练过程,误差反向传播,网络权重参数更新
loss.backward()
optimizer.step()
if batch_idx % cfg.SHOW_STEP == 0:
print('batch %d, loss is %.4f' % (batch_idx, loss.data.item()))
# 误差以及准确率计算
L = 1.0 * len(trainLoader.dataset)
loss = running_loss.data.item() * cfg.BATCH_SIZE / L
accuracy = 100.0 * running_correct.data.item() / L
end_time = time.time()
print(
'epoch %d , train loss is %5.4f and train accuracy is %d/%d = %.2f, consume time is %0.2f min'
% (epoch, loss, running_correct, L, accuracy,
(end_time - start_time) / 60))
return accuracy, loss
def valid_model(validLoader, model, epoch, criterion, cfg):
model.eval()
# 损失及正确率累加变量初始化
running_loss = 0.0
running_correct = 0
start_time = time.time()
with torch.no_grad():
for batch_idx, (data, target) in enumerate(validLoader):
# 针对Cuda进行设置
data, target = data.cuda(), target.cuda()
# 模型前向传播
output = model(data)
# 得到训练误差
loss = criterion(output, target)
# 记录误差及准确率
running_loss += loss.data
# 记录误差及准确率
pred = output.data.max(1, keepdim=True)[1] # 获得得分最高的类别
running_correct += pred.eq(target.data.view_as(pred)).cpu().sum()
if batch_idx % cfg.SHOW_STEP == 0:
print('batch %d, loss is %.4f' % (batch_idx, loss.data.item()))
# 误差以及准确率计算
L = 1.0 * len(validLoader.dataset)
loss = running_loss.data.item() * cfg.BATCH_SIZE / L
accuracy = 100.0 * running_correct.data.item() / L
end_time = time.time()
print(
'epoch %d , test loss is %5.4f and test accuracy is %d/%d = %.2f, time consume is %.2f min'
% (epoch, loss, running_correct, L, accuracy,
(end_time - start_time) / 60))
return accuracy, loss