-
Notifications
You must be signed in to change notification settings - Fork 388
/
train.py
120 lines (100 loc) · 4.3 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
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
from __future__ import print_function
import os
from data import Dataset
import torch
from torch.utils import data
import torch.nn.functional as F
from models import *
import torchvision
from utils import Visualizer, view_model
import torch
import numpy as np
import random
import time
from config import Config
from torch.nn import DataParallel
from torch.optim.lr_scheduler import StepLR
from test import *
def save_model(model, save_path, name, iter_cnt):
save_name = os.path.join(save_path, name + '_' + str(iter_cnt) + '.pth')
torch.save(model.state_dict(), save_name)
return save_name
if __name__ == '__main__':
opt = Config()
if opt.display:
visualizer = Visualizer()
device = torch.device("cuda")
train_dataset = Dataset(opt.train_root, opt.train_list, phase='train', input_shape=opt.input_shape)
trainloader = data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
identity_list = get_lfw_list(opt.lfw_test_list)
img_paths = [os.path.join(opt.lfw_root, each) for each in identity_list]
print('{} train iters per epoch:'.format(len(trainloader)))
if opt.loss == 'focal_loss':
criterion = FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet18':
model = resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet34':
model = resnet34()
elif opt.backbone == 'resnet50':
model = resnet50()
if opt.metric == 'add_margin':
metric_fc = AddMarginProduct(512, opt.num_classes, s=30, m=0.35)
elif opt.metric == 'arc_margin':
metric_fc = ArcMarginProduct(512, opt.num_classes, s=30, m=0.5, easy_margin=opt.easy_margin)
elif opt.metric == 'sphere':
metric_fc = SphereProduct(512, opt.num_classes, m