fast-reid/tools/train.py

94 lines
2.5 KiB
Python
Raw Normal View History

# encoding: utf-8
"""
@author: sherlock
@contact: sherlockliao01@gmail.com
"""
import argparse
import sys
2019-04-21 13:38:55 +08:00
from bisect import bisect_right
2018-10-18 19:04:28 +08:00
from torch.backends import cudnn
2019-01-10 18:39:31 +08:00
sys.path.append('.')
from config import cfg
2019-04-21 13:38:55 +08:00
from data import get_data_bunch
2019-01-10 18:39:31 +08:00
from engine.trainer import do_train
from layers import make_loss
2019-04-21 13:38:55 +08:00
from modeling import build_model
from utils.logger import Logger
from fastai.vision import *
2019-01-10 18:39:31 +08:00
2019-04-23 12:56:38 +08:00
def train(cfg, log_path):
2019-01-10 18:39:31 +08:00
# prepare dataset
2019-04-21 13:38:55 +08:00
data_bunch, test_labels, num_query = get_data_bunch(cfg)
2019-01-10 18:39:31 +08:00
# prepare model
2019-04-21 13:38:55 +08:00
model = build_model(cfg, data_bunch.c)
2019-01-10 18:39:31 +08:00
2019-04-21 13:38:55 +08:00
opt_func = partial(torch.optim.Adam)
2019-01-10 18:39:31 +08:00
2019-04-21 13:38:55 +08:00
def warmup_multistep(start: float, end: float, pct: float) -> float:
warmup_factor = 1
gamma = cfg.SOLVER.GAMMA
milestones = [1.0 * s / cfg.SOLVER.MAX_EPOCHS for s in cfg.SOLVER.STEPS]
warmup_iter = 1.0 * cfg.SOLVER.WARMUP_ITERS / cfg.SOLVER.MAX_EPOCHS
if pct < warmup_iter:
alpha = pct / warmup_iter
warmup_factor = cfg.SOLVER.WARMUP_FACTOR * (1 - alpha) + alpha
return start * warmup_factor * gamma ** bisect_right(milestones, pct)
lr_sched = Scheduler((cfg.SOLVER.BASE_LR, 0), cfg.SOLVER.MAX_EPOCHS, warmup_multistep)
2019-01-10 18:39:31 +08:00
2019-04-21 13:38:55 +08:00
loss_func = make_loss(cfg)
2019-01-10 18:39:31 +08:00
do_train(
cfg,
2019-04-23 12:56:38 +08:00
log_path,
2019-01-10 18:39:31 +08:00
model,
2019-04-21 13:38:55 +08:00
data_bunch,
test_labels,
opt_func,
lr_sched,
2019-01-10 18:39:31 +08:00
loss_func,
num_query
)
def main():
2019-01-10 18:39:31 +08:00
parser = argparse.ArgumentParser(description="ReID Baseline Training")
parser.add_argument(
"--config_file", default="", help="path to config file", type=str
)
parser.add_argument("opts", help="Modify config options using the command-line", default=None,
nargs=argparse.REMAINDER)
args = parser.parse_args()
2019-01-10 18:39:31 +08:00
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
if args.config_file != "":
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
2019-04-23 12:56:38 +08:00
log_path = Path(os.path.join(cfg.OUTPUT_DIR, 'log.txt'))
with log_path.open('w') as f: f.write('{}\n'.format(args))
2019-04-21 13:38:55 +08:00
print(args)
2019-01-10 18:39:31 +08:00
if args.config_file != "":
2019-04-21 13:38:55 +08:00
print("Loaded configuration file {}".format(args.config_file))
2019-01-10 18:39:31 +08:00
with open(args.config_file, 'r') as cf:
config_str = "\n" + cf.read()
2019-04-21 13:38:55 +08:00
print(config_str)
print("Running with config:\n{}".format(cfg))
2019-04-23 12:56:38 +08:00
with log_path.open('a') as f: f.write('{}\n'.format(cfg))
2018-10-18 19:04:28 +08:00
cudnn.benchmark = True
2019-04-23 12:56:38 +08:00
train(cfg, log_path)
if __name__ == '__main__':
main()