def train(cfg):
# Set up environment.
init_distributed_training(cfg)
local_rank_id = get_local_rank()
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED + 10 * local_rank_id)
torch.manual_seed(cfg.RNG_SEED + 10 * local_rank_id)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Setup logging format.
logging.setup_logging(cfg.OUTPUT_DIR)
logger.info('init start')
# 迭代轮数从１开始计数
arguments = {"cur_epoch": 1}
device = get_device(local_rank_id)
model = build_recognizer(cfg, device)
criterion = build_criterion(cfg, device)
optimizer = build_optimizer(cfg, model)
lr_scheduler = build_lr_scheduler(cfg, optimizer)
checkpointer = CheckPointer(model,
optimizer=optimizer,
scheduler=lr_scheduler,
save_dir=cfg.OUTPUT_DIR,
save_to_disk=True)
if cfg.TRAIN.RESUME:
logger.info('resume start')
extra_checkpoint_data = checkpointer.load(map_location=device)
if isinstance(extra_checkpoint_data, dict):
arguments['cur_epoch'] = extra_checkpoint_data['cur_epoch']
if cfg.LR_SCHEDULER.IS_WARMUP:
logger.info('warmup start')
if lr_scheduler.finished:
optimizer.load_state_dict(
lr_scheduler.after_scheduler.optimizer.state_dict())
else:
optimizer.load_state_dict(
lr_scheduler.optimizer.state_dict())
lr_scheduler.optimizer = optimizer
lr_scheduler.after_scheduler.optimizer = optimizer
logger.info('warmup end')
logger.info('resume end')
data_loader = build_dataloader(cfg, is_train=True)
logger.info('init end')
synchronize()
do_train(cfg, arguments, data_loader, model, criterion, optimizer,
lr_scheduler, checkpointer, device)
def test(cfg):
# Set up environment.
init_distributed_training(cfg)
# Set random seed from configs.
np.random.seed(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
device = get_device(local_rank=get_local_rank())
model = build_recognizer(cfg, device=device)
synchronize()
do_evaluation(cfg, model, device)
def main(data_shape, config_file, mobile_name):
cfg.merge_from_file(config_file)
device = get_device(local_rank=get_local_rank())
model = build_recognizer(cfg, device)
data = torch.randn(data_shape).to(device=device, non_blocking=True)
GFlops, params_size = compute_num_flops(model, data)
print(f'{mobile_name} ' + '*' * 10)
print(f'device: {device}')
print(f'GFlops: {GFlops}')
print(f'Params Size: {params_size}')
total_time = 0.0
num = 100
for i in range(num):
data = torch.randn(data_shape).to(device=device, non_blocking=True)
start = time.time()
model(data)
total_time = time.time() - start
print(f'one process need {total_time / num}')
def main(data_shape, config_file, mobile_name):
cfg.merge_from_file(config_file)
device = get_device(local_rank=get_local_rank())
model = build_recognizer(cfg, device)
model.eval()
data = torch.randn(data_shape).to(device=device, non_blocking=True)
GFlops, params_size = compute_num_flops(model, data)
print(f'{mobile_name} ' + '*' * 10)
print(f'device: {device}')
print(f'GFlops: {GFlops}')
print(f'Params Size: {params_size}')
data = torch.randn(data_shape)
t1 = 0.0
num = 100
begin = time.time()
for i in range(num):
start = time.time()
model(data.to(device=device, non_blocking=True))
t1 += time.time() - start
t2 = time.time() - begin
print(f'one process need {t2 / num}, model compute need: {t1 / num}')