-
Notifications
You must be signed in to change notification settings - Fork 0
/
lr_optimizer.py
109 lines (91 loc) · 4.28 KB
/
lr_optimizer.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
import matplotlib.pyplot as plt
import os
import numpy as np
import json
import argparse
import torch
from torch.optim import lr_scheduler
from tqdm import tqdm
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 trainer import Trainer
from utils import Logger
from utils import util
from utils import torchsummary
from utils import viewTraining
from utils import lr_finder
from utils import classActivationMap
import importlib
import math
import torchvision
from torch.nn import functional as F
from torch import topk
import skimage.transform
print("Modules loaded")
def get_instance(module, name, config, *args):
return getattr(module, config[name]['type'])(*args, **config[name]['args'])
def main(config, resume):
#torch.set_num_threads(16)
#print("Using 16 of " + str(torch.get_num_threads()) + " threads")
print("GPUs available: " + str(torch.cuda.device_count()))
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
os.environ["NUMEXPR_MAX_THREADS"] = "16"
os.environ["OMP_NUM_THREADS"] = "16"
train_logger = Logger()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# setup data_loader instances
data_loader = get_instance(module_data, 'data_loader', config)
valid_data_loader = data_loader.split_validation()
# build model architecture
model = get_instance(module_arch, 'arch', config)
print(model)
#torchsummary.summary(model, (1,7,32,32))
if torch.cuda.is_available():
for gp in range(torch.cuda.device_count()):
print("Using GPU " + str(gp+1) + "/" + str(torch.cuda.device_count()) + ": " + torch.cuda.get_device_name(gp))
else:
print("Using CPU to train")
# get function handles of loss and metrics
loss = getattr(module_loss, config['loss']) #looks in model/loss.py for criterion function specified in config
criterion = loss(data_loader.dataset.weight.to(device)) # for imbalanced datasets
metrics = [getattr(module_metric, met) for met in config['metrics']]
# build optimizer, learning rate scheduler. delete every lines containing lr_scheduler for disabling scheduler
trainable_params = filter(lambda p: p.requires_grad, model.parameters())
optimizer = get_instance(torch.optim, 'optimizer', config, trainable_params)
lr_scheduler = get_instance(torch.optim.lr_scheduler, 'lr_scheduler', config, optimizer)
trainer = Trainer(model, criterion, metrics, optimizer,
resume=resume,
config=config,
data_loader=data_loader,
valid_data_loader=valid_data_loader,
lr_scheduler=lr_scheduler,
train_logger=train_logger)
lr_finder_train = lr_finder.LRFinder(model, optimizer, criterion, device = device)
lr_finder_train.range_test(data_loader, end_lr=0.5, num_iter = 100, val_loader = valid_data_loader, num_epochs = 5)
lr_finder_train.plot(skip_start = 10, skip_end=5)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='PyTorch Template')
parser.add_argument('-c', '--config', default=None, type=str,
help='config file path (default: None)')
parser.add_argument('-r', '--resume', default=None, type=str,
help='path to latest checkpoint (default: None)')
parser.add_argument('-d', '--device', default=None, type=str,
help='indices of GPUs to enable (default: all)')
args = parser.parse_args()
if args.config:
# load config file
with open(args.config) as handle:
config = json.load(handle)
# setting path to save trained models and log files
path = os.path.join(config['trainer']['save_dir'], config['name'])
elif args.resume:
# load config from checkpoint if new config file is not given.
# Use '--config' and '--resume' together to fine-tune trained model with changed configurations.
config = torch.load(args.resume)['config']
else:
raise AssertionError("Configuration file need to be specified. Add '-c config.json', for example.")
if args.device:
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
main(config, args.resume)