changed main to main_test
Huy Huynh
parent
8d7688eab2
commit
3e9c55ae28
191
scripts/main.py
191
scripts/main.py
|
|
@ -1,191 +0,0 @@
|
||||||
import sys
|
|
||||||
import time
|
|
||||||
import os.path as osp
|
|
||||||
import argparse
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
|
|
||||||
import torchreid
|
|
||||||
from torchreid.utils import (
|
|
||||||
Logger, check_isfile, set_random_seed, collect_env_info,
|
|
||||||
resume_from_checkpoint, load_pretrained_weights, compute_model_complexity
|
|
||||||
)
|
|
||||||
|
|
||||||
from default_config import (
|
|
||||||
imagedata_kwargs, optimizer_kwargs, videodata_kwargs, engine_run_kwargs,
|
|
||||||
get_default_config, lr_scheduler_kwargs
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def build_datamanager(cfg):
|
|
||||||
if cfg.data.type == 'image':
|
|
||||||
return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
|
|
||||||
else:
|
|
||||||
return torchreid.data.VideoDataManager(**videodata_kwargs(cfg))
|
|
||||||
|
|
||||||
|
|
||||||
def build_engine(cfg, datamanager, model, optimizer, scheduler):
|
|
||||||
if cfg.data.type == 'image':
|
|
||||||
if cfg.loss.name == 'softmax':
|
|
||||||
engine = torchreid.engine.ImageSoftmaxEngine(
|
|
||||||
datamanager,
|
|
||||||
model,
|
|
||||||
optimizer=optimizer,
|
|
||||||
scheduler=scheduler,
|
|
||||||
use_gpu=cfg.use_gpu,
|
|
||||||
label_smooth=cfg.loss.softmax.label_smooth
|
|
||||||
)
|
|
||||||
|
|
||||||
else:
|
|
||||||
engine = torchreid.engine.ImageTripletEngine(
|
|
||||||
datamanager,
|
|
||||||
model,
|
|
||||||
optimizer=optimizer,
|
|
||||||
margin=cfg.loss.triplet.margin,
|
|
||||||
weight_t=cfg.loss.triplet.weight_t,
|
|
||||||
weight_x=cfg.loss.triplet.weight_x,
|
|
||||||
scheduler=scheduler,
|
|
||||||
use_gpu=cfg.use_gpu,
|
|
||||||
label_smooth=cfg.loss.softmax.label_smooth
|
|
||||||
)
|
|
||||||
|
|
||||||
else:
|
|
||||||
if cfg.loss.name == 'softmax':
|
|
||||||
engine = torchreid.engine.VideoSoftmaxEngine(
|
|
||||||
datamanager,
|
|
||||||
model,
|
|
||||||
optimizer=optimizer,
|
|
||||||
scheduler=scheduler,
|
|
||||||
use_gpu=cfg.use_gpu,
|
|
||||||
label_smooth=cfg.loss.softmax.label_smooth,
|
|
||||||
pooling_method=cfg.video.pooling_method
|
|
||||||
)
|
|
||||||
|
|
||||||
else:
|
|
||||||
engine = torchreid.engine.VideoTripletEngine(
|
|
||||||
datamanager,
|
|
||||||
model,
|
|
||||||
optimizer=optimizer,
|
|
||||||
margin=cfg.loss.triplet.margin,
|
|
||||||
weight_t=cfg.loss.triplet.weight_t,
|
|
||||||
weight_x=cfg.loss.triplet.weight_x,
|
|
||||||
scheduler=scheduler,
|
|
||||||
use_gpu=cfg.use_gpu,
|
|
||||||
label_smooth=cfg.loss.softmax.label_smooth
|
|
||||||
)
|
|
||||||
|
|
||||||
return engine
|
|
||||||
|
|
||||||
|
|
||||||
def reset_config(cfg, args):
|
|
||||||
if args.root:
|
|
||||||
cfg.data.root = args.root
|
|
||||||
if args.sources:
|
|
||||||
cfg.data.sources = args.sources
|
|
||||||
if args.targets:
|
|
||||||
cfg.data.targets = args.targets
|
|
||||||
if args.transforms:
|
|
||||||
cfg.data.transforms = args.transforms
|
|
||||||
|
|
||||||
|
|
||||||
def check_cfg(cfg):
|
|
||||||
if cfg.loss.name == 'triplet' and cfg.loss.triplet.weight_x == 0:
|
|
||||||
assert cfg.train.fixbase_epoch == 0, \
|
|
||||||
'The output of classifier is not included in the computational graph'
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'--config-file', type=str, default='', help='path to config file'
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'-s',
|
|
||||||
'--sources',
|
|
||||||
type=str,
|
|
||||||
nargs='+',
|
|
||||||
help='source datasets (delimited by space)'
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'-t',
|
|
||||||
'--targets',
|
|
||||||
type=str,
|
|
||||||
nargs='+',
|
|
||||||
help='target datasets (delimited by space)'
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'--transforms', type=str, nargs='+', help='data augmentation'
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'--root', type=str, default='', help='path to data root'
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
'opts',
|
|
||||||
default=None,
|
|
||||||
nargs=argparse.REMAINDER,
|
|
||||||
help='Modify config options using the command-line'
|
|
||||||
)
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
cfg = get_default_config()
|
|
||||||
cfg.use_gpu = torch.cuda.is_available()
|
|
||||||
if args.config_file:
|
|
||||||
cfg.merge_from_file(args.config_file)
|
|
||||||
reset_config(cfg, args)
|
|
||||||
cfg.merge_from_list(args.opts)
|
|
||||||
set_random_seed(cfg.train.seed)
|
|
||||||
check_cfg(cfg)
|
|
||||||
|
|
||||||
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
|
|
||||||
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
|
|
||||||
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
|
|
||||||
|
|
||||||
print('Show configuration\n{}\n'.format(cfg))
|
|
||||||
print('Collecting env info ...')
|
|
||||||
print('** System info **\n{}\n'.format(collect_env_info()))
|
|
||||||
|
|
||||||
if cfg.use_gpu:
|
|
||||||
torch.backends.cudnn.benchmark = True
|
|
||||||
|
|
||||||
datamanager = build_datamanager(cfg)
|
|
||||||
|
|
||||||
print('Building model: {}'.format(cfg.model.name))
|
|
||||||
model = torchreid.models.build_model(
|
|
||||||
name=cfg.model.name,
|
|
||||||
num_classes=datamanager.num_train_pids,
|
|
||||||
loss=cfg.loss.name,
|
|
||||||
pretrained=cfg.model.pretrained,
|
|
||||||
use_gpu=cfg.use_gpu
|
|
||||||
)
|
|
||||||
num_params, flops = compute_model_complexity(
|
|
||||||
model, (1, 3, cfg.data.height, cfg.data.width)
|
|
||||||
)
|
|
||||||
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
|
|
||||||
|
|
||||||
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
|
|
||||||
load_pretrained_weights(model, cfg.model.load_weights)
|
|
||||||
|
|
||||||
if cfg.use_gpu:
|
|
||||||
model = nn.DataParallel(model).cuda()
|
|
||||||
|
|
||||||
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
|
|
||||||
scheduler = torchreid.optim.build_lr_scheduler(
|
|
||||||
optimizer, **lr_scheduler_kwargs(cfg)
|
|
||||||
)
|
|
||||||
|
|
||||||
if cfg.model.resume and check_isfile(cfg.model.resume):
|
|
||||||
cfg.train.start_epoch = resume_from_checkpoint(
|
|
||||||
cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler
|
|
||||||
)
|
|
||||||
|
|
||||||
print(
|
|
||||||
'Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)
|
|
||||||
)
|
|
||||||
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
|
|
||||||
engine.run(**engine_run_kwargs(cfg))
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
main()
|
|
||||||
|
|
@ -1,49 +0,0 @@
|
||||||
if __name__ == '__main__':
|
|
||||||
import torchreid
|
|
||||||
|
|
||||||
datamanager = torchreid.data.ImageDataManager(
|
|
||||||
root='reid-data',
|
|
||||||
sources='market1501',
|
|
||||||
targets='market1501',
|
|
||||||
height=256,
|
|
||||||
width=128,
|
|
||||||
batch_size_train=32,
|
|
||||||
batch_size_test=100,
|
|
||||||
transforms=['random_flip', 'color_jitter']
|
|
||||||
)
|
|
||||||
|
|
||||||
model = torchreid.models.build_model(
|
|
||||||
name='osnet_x1_0',
|
|
||||||
num_classes=datamanager.num_train_pids,
|
|
||||||
loss='softmax',
|
|
||||||
pretrained=True
|
|
||||||
)
|
|
||||||
model = model.cuda()
|
|
||||||
|
|
||||||
weight_path = 'log\osnet_x1_0_market_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip.pth'
|
|
||||||
torchreid.utils.load_pretrained_weights(model, weight_path)
|
|
||||||
|
|
||||||
optimizer = torchreid.optim.build_optimizer(model, optim='adam', lr=0.0003)
|
|
||||||
|
|
||||||
|
|
||||||
# scheduler = torchreid.optim.build_lr_scheduler(
|
|
||||||
# optimizer, lr_scheduler='single_step', stepsize=20
|
|
||||||
# )
|
|
||||||
|
|
||||||
|
|
||||||
engine = torchreid.engine.ImageSoftmaxEngine(
|
|
||||||
datamanager,
|
|
||||||
model,
|
|
||||||
optimizer=optimizer,
|
|
||||||
# scheduler=scheduler,
|
|
||||||
label_smooth=True
|
|
||||||
)
|
|
||||||
|
|
||||||
engine.run(
|
|
||||||
save_dir='log/osnet_ibn_x1_0',
|
|
||||||
# max_epoch=60,
|
|
||||||
# eval_freq=10,
|
|
||||||
# print_freq=10,
|
|
||||||
test_only=True,
|
|
||||||
visrank=True
|
|
||||||
)
|
|
||||||
Loading…
Reference in New Issue