-
Notifications
You must be signed in to change notification settings - Fork 1
/
train.py
115 lines (100 loc) · 5.4 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
import argparse
import collections
import pprint
import torch
import numpy as np
import data_loader.data_loaders as module_data
import model.loss as module_loss
import model.metric as module_metric
import model.model as module_arch
from parse_config import ConfigParser
from trainer import Trainer
deterministic = False
if deterministic:
# fix random seeds for reproducibility
SEED = 123
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(SEED)
def learing_rate_scheduler(optimizer, config):
if "type" in config._config["lr_scheduler"]:
if config["lr_scheduler"]["type"] == "CustomLR": # linear learning rate decay
lr_scheduler_args = config["lr_scheduler"]["args"]
gamma = lr_scheduler_args["gamma"] if "gamma" in lr_scheduler_args else 0.1
print("Scheduler step1, step2, warmup_epoch, gamma:", (lr_scheduler_args["step1"], lr_scheduler_args["step2"], lr_scheduler_args["warmup_epoch"], gamma))
def lr_lambda(epoch):
if lr_scheduler_args.get('step3') is not None and epoch >= lr_scheduler_args['step3']:
lr = gamma * gamma * gamma
elif epoch >= lr_scheduler_args.get("step2", float('inf')): # 默认值设为无穷大，确保跳过
lr = gamma * gamma
elif epoch >= lr_scheduler_args.get("step1", float('inf')):
lr = gamma
else:
lr = 1
"""Warmup"""
warmup_epoch = lr_scheduler_args["warmup_epoch"]
if epoch < warmup_epoch:
lr = lr * float(1 + epoch) / warmup_epoch
return lr
lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
else:
# lr_scheduler = config.init_obj('lr_scheduler', torch.optim.lr_scheduler, optimizer) # cosine learning rate decay
# lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60, 120, 160], gamma=0.2) #learning rate decay
else:
lr_scheduler = None
return lr_scheduler
def main(config):
logger = config.get_logger('train')
# setup data_loader instances
data_loader = config.init_obj('data_loader', module_data)
valid_data_loader = data_loader.split_validation()
# build model architecture, then print to console
model = config.init_obj('arch', module_arch)
logger.info(model)
# get function handles of loss and metrics
loss_class = getattr(module_loss, config["loss"]["type"])
if hasattr(loss_class, "require_num_experts") and loss_class.require_num_experts:
criterion = config.init_obj('loss', module_loss, cls_num_list=data_loader.cls_num_list, num_experts=config["arch"]["args"]["num_experts"])
else:
criterion = config.init_obj('loss', module_loss, cls_num_list=data_loader.cls_num_list)
metrics = [getattr(module_metric, met) for met in config['metrics']]
# build optimizer, learning rate scheduler.
optimizer = config.init_obj('optimizer', torch.optim, model.parameters())
lr_scheduler = learing_rate_scheduler(optimizer, config)
trainer = Trainer(model, criterion, metrics, optimizer,
config=config,
data_loader=data_loader,
valid_data_loader=valid_data_loader,
lr_scheduler=lr_scheduler)
trainer.train()
if __name__ == '__main__':
args = argparse.ArgumentParser(description='PyTorch Template')
args.add_argument('-c', '--config', default=None, type=str,
help='config file path (default: None)')
args.add_argument('-r', '--resume', default=None, type=str,
help='path to latest checkpoint (default: None)')
args.add_argument('-d', '--device', default=None, type=str,
help='indices of GPUs to enable (default: all)')
# custom cli options to modify configuration from default values given in json file.
CustomArgs = collections.namedtuple('CustomArgs', 'flags type target')
options = [
CustomArgs(['--lr', '--learning_rate'], type=float, target='optimizer;args;lr'),
CustomArgs(['--bs', '--batch_size'], type=int, target='data_loader;args;batch_size'),
CustomArgs(['--name'], type=str, target='name'),
CustomArgs(['--epochs'], type=int, target='trainer;epochs'),
CustomArgs(['--step1'], type=int, target='lr_scheduler;args;step1'),
CustomArgs(['--step2'], type=int, target='lr_scheduler;args;step2'),
CustomArgs(['--warmup'], type=int, target='lr_scheduler;args;warmup_epoch'),
CustomArgs(['--gamma'], type=float, target='lr_scheduler;args;gamma'),
CustomArgs(['--save_period'], type=int, target='trainer;save_period'),
CustomArgs(['--reduce_dimension'], type=int, target='arch;args;reduce_dimension'),
CustomArgs(['--layer2_dimension'], type=int, target='arch;args;layer2_output_dim'),
CustomArgs(['--layer3_dimension'], type=int, target='arch;args;layer3_output_dim'),
CustomArgs(['--layer4_dimension'], type=int, target='arch;args;layer4_output_dim'),
CustomArgs(['--num_experts'], type=int, target='arch;args;num_experts')
]
config = ConfigParser.from_args(args, options)
pprint.pprint(config)
main(config)