def get_lr_scheduler(cfg, optimizer):
it_ep = num_iters_per_epoch(cfg) #一次训练的最大读取次数(781)
if cfg.linear_final_lr is None and cfg.cosine_minimum is None:
lr_iter_boundaries = [it_ep * ep for ep in cfg.lr_epoch_boundaries
] #lr_epoch_boundaries=[200, 400]
return WarmupMultiStepLR(
optimizer,
lr_iter_boundaries,
cfg.lr_decay_factor,
warmup_factor=cfg.warmup_factor,
warmup_iters=cfg.warmup_epochs * it_ep,
warmup_method=cfg.warmup_method,
)
elif cfg.cosine_minimum is None:
return WarmupLinearLR(
optimizer,
final_lr=cfg.linear_final_lr,
final_iters=cfg.max_epochs * it_ep,
warmup_factor=cfg.warmup_factor,
warmup_iters=cfg.warmup_epochs * it_ep,
warmup_method=cfg.warmup_method,
)
else:
assert cfg.warmup_epochs == 0
assert cfg.linear_final_lr is None
assert cfg.lr_decay_factor is None
if cfg.lr_epoch_boundaries is None:
print('use cosine decay, the minimum is ', cfg.cosine_minimum)
return CosineAnnealingLR(optimizer=optimizer,
T_max=cfg.max_epochs * it_ep,
eta_min=cfg.cosine_minimum)
else:
assert len(cfg.lr_epoch_boundaries) == 1
assert cfg.cosine_minimum > 0
print('use extended cosine decay, the minimum is ',
cfg.cosine_minimum)
return CosineAnnealingExtendLR(
optimizer=optimizer,
T_cosine_max=cfg.lr_epoch_boundaries[0] * it_ep,
eta_min=cfg.cosine_minimum)
def csgd_train_main(local_rank,
cfg: BaseConfigByEpoch,
target_deps,
succeeding_strategy,
pacesetter_dict,
centri_strength,
pruned_weights,
net=None,
train_dataloader=None,
val_dataloader=None,
show_variables=False,
convbuilder=None,
init_hdf5=None,
no_l2_keywords='depth',
use_nesterov=False,
load_weights_keyword=None,
keyword_to_lr_mult=None,
auto_continue=False,
save_hdf5_epochs=10000):
ensure_dir(cfg.output_dir)
ensure_dir(cfg.tb_dir)
clusters_save_path = os.path.join(cfg.output_dir, 'clusters.npy')
with Engine(local_rank=local_rank) as engine:
engine.setup_log(name='train',
log_dir=cfg.output_dir,
file_name='log.txt')
# ----------------------------- build model ------------------------------
if convbuilder is None:
convbuilder = ConvBuilder(base_config=cfg)
if net is None:
net_fn = get_model_fn(cfg.dataset_name, cfg.network_type)
model = net_fn(cfg, convbuilder)
else:
model = net
model = model.cuda()
# ----------------------------- model done ------------------------------
# ---------------------------- prepare data -------------------------
if train_dataloader is None:
train_data = create_dataset(cfg.dataset_name,
cfg.dataset_subset,
cfg.global_batch_size,
distributed=engine.distributed)
if cfg.val_epoch_period > 0 and val_dataloader is None:
val_data = create_dataset(cfg.dataset_name,
'val',
global_batch_size=100,
distributed=False)
engine.echo('NOTE: Data prepared')
engine.echo(
'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}'
.format(cfg.global_batch_size, torch.cuda.device_count(),
torch.cuda.memory_allocated()))
# ----------------------------- data done --------------------------------
# ------------------------ parepare optimizer, scheduler, criterion -------
if no_l2_keywords is None:
no_l2_keywords = []
if type(no_l2_keywords) is not list:
no_l2_keywords = [no_l2_keywords]
# For a target parameter, cancel its weight decay in optimizer, because the weight decay will be later encoded in the decay mat
conv_idx = 0
for k, v in model.named_parameters():
if v.dim() != 4:
continue
print('prune {} from {} to {}'.format(conv_idx,
target_deps[conv_idx],
cfg.deps[conv_idx]))
if target_deps[conv_idx] < cfg.deps[conv_idx]:
no_l2_keywords.append(k.replace(KERNEL_KEYWORD, 'conv'))
no_l2_keywords.append(k.replace(KERNEL_KEYWORD, 'bn'))
conv_idx += 1
print('no l2: ', no_l2_keywords)
optimizer = get_optimizer(engine,
cfg,
model,
no_l2_keywords=no_l2_keywords,
use_nesterov=use_nesterov,
keyword_to_lr_mult=keyword_to_lr_mult)
scheduler = get_lr_scheduler(cfg, optimizer)
criterion = get_criterion(cfg).cuda()
# --------------------------------- done -------------------------------
engine.register_state(scheduler=scheduler,
model=model,
optimizer=optimizer)
if engine.distributed:
torch.cuda.set_device(local_rank)
engine.echo('Distributed training, device {}'.format(local_rank))
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
broadcast_buffers=False,
)
else:
assert torch.cuda.device_count() == 1
engine.echo('Single GPU training')
if cfg.init_weights:
engine.load_checkpoint(cfg.init_weights)
if init_hdf5:
engine.load_hdf5(init_hdf5,
load_weights_keyword=load_weights_keyword)
if auto_continue:
assert cfg.init_weights is None
engine.load_checkpoint(get_last_checkpoint(cfg.output_dir))
if show_variables:
engine.show_variables()
# ===================================== prepare the clusters and matrices for C-SGD ==========
kernel_namedvalue_list = engine.get_all_conv_kernel_namedvalue_as_list(
)
if os.path.exists(clusters_save_path):
layer_idx_to_clusters = np.load(clusters_save_path,
allow_pickle=True).item()
else:
if local_rank == 0:
layer_idx_to_clusters = get_layer_idx_to_clusters(
kernel_namedvalue_list=kernel_namedvalue_list,
target_deps=target_deps,
pacesetter_dict=pacesetter_dict)
if pacesetter_dict is not None:
for follower_idx, pacesetter_idx in pacesetter_dict.items(
):
if pacesetter_idx in layer_idx_to_clusters:
layer_idx_to_clusters[
follower_idx] = layer_idx_to_clusters[
pacesetter_idx]
np.save(clusters_save_path, layer_idx_to_clusters)
else:
while not os.path.exists(clusters_save_path):
time.sleep(10)
print('sleep, waiting for process 0 to calculate clusters')
layer_idx_to_clusters = np.load(clusters_save_path,
allow_pickle=True).item()
param_name_to_merge_matrix = generate_merge_matrix_for_kernel(
deps=cfg.deps,
layer_idx_to_clusters=layer_idx_to_clusters,
kernel_namedvalue_list=kernel_namedvalue_list)
add_vecs_to_merge_mat_dicts(param_name_to_merge_matrix)
param_name_to_decay_matrix = generate_decay_matrix_for_kernel_and_vecs(
deps=cfg.deps,
layer_idx_to_clusters=layer_idx_to_clusters,
kernel_namedvalue_list=kernel_namedvalue_list,
weight_decay=cfg.weight_decay,
weight_decay_bias=cfg.weight_decay_bias,
centri_strength=centri_strength)
print(param_name_to_decay_matrix.keys())
print(param_name_to_merge_matrix.keys())
conv_idx = 0
param_to_clusters = {}
for k, v in model.named_parameters():
if v.dim() != 4:
continue
if conv_idx in layer_idx_to_clusters:
for clsts in layer_idx_to_clusters[conv_idx]:
if len(clsts) > 1:
param_to_clusters[v] = layer_idx_to_clusters[conv_idx]
break
conv_idx += 1
# ============================================================================================
# ------------ do training ---------------------------- #
engine.log("\n\nStart training with pytorch version {}".format(
torch.__version__))
iteration = engine.state.iteration
iters_per_epoch = num_iters_per_epoch(cfg)
max_iters = iters_per_epoch * cfg.max_epochs
tb_writer = SummaryWriter(cfg.tb_dir)
tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss']
model.train()
done_epochs = iteration // iters_per_epoch
last_epoch_done_iters = iteration % iters_per_epoch
if done_epochs == 0 and last_epoch_done_iters == 0:
engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5'))
recorded_train_time = 0
recorded_train_examples = 0
collected_train_loss_sum = 0
collected_train_loss_count = 0
for epoch in range(done_epochs, cfg.max_epochs):
if engine.distributed and hasattr(train_data, 'train_sampler'):
train_data.train_sampler.set_epoch(epoch)
if epoch == done_epochs:
pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters))
else:
pbar = tqdm(range(iters_per_epoch))
if epoch == 0 and local_rank == 0:
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str='Init',
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
top1 = AvgMeter()
top5 = AvgMeter()
losses = AvgMeter()
discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs)
pbar.set_description('Train' + discrip_str)
for _ in pbar:
start_time = time.time()
data, label = load_cuda_data(train_data,
dataset_name=cfg.dataset_name)
# load_cuda_data(train_dataloader, cfg.dataset_name)
data_time = time.time() - start_time
train_net_time_start = time.time()
acc, acc5, loss = train_one_step(
model,
data,
label,
optimizer,
criterion,
param_name_to_merge_matrix=param_name_to_merge_matrix,
param_name_to_decay_matrix=param_name_to_decay_matrix)
train_net_time_end = time.time()
if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters:
recorded_train_examples += cfg.global_batch_size
recorded_train_time += train_net_time_end - train_net_time_start
scheduler.step()
for module in model.modules():
if hasattr(module, 'set_cur_iter'):
module.set_cur_iter(iteration)
if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0:
for tag, value in zip(
tb_tags,
[acc.item(), acc5.item(),
loss.item()]):
tb_writer.add_scalars(tag, {'Train': value}, iteration)
deviation_sum = 0
for param, clusters in param_to_clusters.items():
pvalue = param.detach().cpu().numpy()
for cl in clusters:
if len(cl) == 1:
continue
selected = pvalue[cl, :, :, :]
mean_kernel = np.mean(selected,
axis=0,
keepdims=True)
diff = selected - mean_kernel
deviation_sum += np.sum(diff**2)
tb_writer.add_scalars('deviation_sum',
{'Train': deviation_sum}, iteration)
top1.update(acc.item())
top5.update(acc5.item())
losses.update(loss.item())
if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS:
collected_train_loss_sum += loss.item()
collected_train_loss_count += 1
pbar_dic = OrderedDict()
pbar_dic['data-time'] = '{:.2f}'.format(data_time)
pbar_dic['cur_iter'] = iteration
pbar_dic['lr'] = scheduler.get_lr()[0]
pbar_dic['top1'] = '{:.5f}'.format(top1.mean)
pbar_dic['top5'] = '{:.5f}'.format(top5.mean)
pbar_dic['loss'] = '{:.5f}'.format(losses.mean)
pbar.set_postfix(pbar_dic)
iteration += 1
if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0:
engine.update_iteration(iteration)
if (not engine.distributed) or (engine.distributed and
engine.world_rank == 0):
engine.save_and_link_checkpoint(cfg.output_dir)
if iteration >= max_iters:
break
# do something after an epoch?
engine.update_iteration(iteration)
engine.save_latest_ckpt(cfg.output_dir)
if (epoch + 1) % save_hdf5_epochs == 0:
engine.save_hdf5(
os.path.join(cfg.output_dir,
'epoch-{}.hdf5'.format(epoch)))
if local_rank == 0 and \
cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0):
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str=discrip_str,
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
if iteration >= max_iters:
break
# do something after the training
if recorded_train_time > 0:
exp_per_sec = recorded_train_examples / recorded_train_time
else:
exp_per_sec = 0
engine.log(
'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}'
.format(int(TRAIN_SPEED_START * max_iters),
int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size,
recorded_train_examples, recorded_train_time, exp_per_sec))
if cfg.save_weights:
engine.save_checkpoint(cfg.save_weights)
print('NOTE: training finished, saved to {}'.format(
cfg.save_weights))
engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5'))
if collected_train_loss_count > 0:
engine.log(
'TRAIN LOSS collected over last {} epochs: {:.6f}'.format(
COLLECT_TRAIN_LOSS_EPOCHS,
collected_train_loss_sum / collected_train_loss_count))
if local_rank == 0:
csgd_prune_and_save(engine=engine,
layer_idx_to_clusters=layer_idx_to_clusters,
save_file=pruned_weights,
succeeding_strategy=succeeding_strategy,
new_deps=target_deps)
def train_main(local_rank,
cfg: BaseConfigByEpoch,
net=None,
train_dataloader=None,
val_dataloader=None,
show_variables=False,
convbuilder=None,
init_hdf5=None,
no_l2_keywords='depth',
gradient_mask=None,
use_nesterov=False,
tensorflow_style_init=False,
load_weights_keyword=None,
keyword_to_lr_mult=None,
auto_continue=False,
lasso_keyword_to_strength=None,
save_hdf5_epochs=10000):
if no_l2_keywords is None:
no_l2_keywords = []
if type(no_l2_keywords) is not list:
no_l2_keywords = [no_l2_keywords]
ensure_dir(cfg.output_dir)
ensure_dir(cfg.tb_dir)
with Engine(local_rank=local_rank) as engine:
engine.setup_log(name='train',
log_dir=cfg.output_dir,
file_name='log.txt')
# ----------------------------- build model ------------------------------
if convbuilder is None:
convbuilder = ConvBuilder(base_config=cfg)
if net is None:
net_fn = get_model_fn(cfg.dataset_name, cfg.network_type)
model = net_fn(cfg, convbuilder)
else:
model = net
model = model.cuda()
# ----------------------------- model done ------------------------------
# ---------------------------- prepare data -------------------------
if train_dataloader is None:
train_data = create_dataset(cfg.dataset_name,
cfg.dataset_subset,
cfg.global_batch_size,
distributed=engine.distributed)
if cfg.val_epoch_period > 0 and val_dataloader is None:
val_data = create_dataset(cfg.dataset_name,
'val',
global_batch_size=100,
distributed=False)
engine.echo('NOTE: Data prepared')
engine.echo(
'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}'
.format(cfg.global_batch_size, torch.cuda.device_count(),
torch.cuda.memory_allocated()))
# ----------------------------- data done --------------------------------
# ------------------------ parepare optimizer, scheduler, criterion -------
optimizer = get_optimizer(engine,
cfg,
model,
no_l2_keywords=no_l2_keywords,
use_nesterov=use_nesterov,
keyword_to_lr_mult=keyword_to_lr_mult)
scheduler = get_lr_scheduler(cfg, optimizer)
criterion = get_criterion(cfg).cuda()
# --------------------------------- done -------------------------------
engine.register_state(scheduler=scheduler,
model=model,
optimizer=optimizer)
if engine.distributed:
torch.cuda.set_device(local_rank)
engine.echo('Distributed training, device {}'.format(local_rank))
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
broadcast_buffers=False,
)
else:
assert torch.cuda.device_count() == 1
engine.echo('Single GPU training')
if tensorflow_style_init:
init_as_tensorflow(model)
if cfg.init_weights:
engine.load_checkpoint(cfg.init_weights)
if init_hdf5:
engine.load_hdf5(init_hdf5,
load_weights_keyword=load_weights_keyword)
if auto_continue:
assert cfg.init_weights is None
engine.load_checkpoint(get_last_checkpoint(cfg.output_dir))
if show_variables:
engine.show_variables()
# ------------ do training ---------------------------- #
engine.log("\n\nStart training with pytorch version {}".format(
torch.__version__))
iteration = engine.state.iteration
iters_per_epoch = num_iters_per_epoch(cfg)
max_iters = iters_per_epoch * cfg.max_epochs
tb_writer = SummaryWriter(cfg.tb_dir)
tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss']
model.train()
done_epochs = iteration // iters_per_epoch
last_epoch_done_iters = iteration % iters_per_epoch
if done_epochs == 0 and last_epoch_done_iters == 0:
engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5'))
recorded_train_time = 0
recorded_train_examples = 0
collected_train_loss_sum = 0
collected_train_loss_count = 0
if gradient_mask is not None:
gradient_mask_tensor = {}
for name, value in gradient_mask.items():
gradient_mask_tensor[name] = torch.Tensor(value).cuda()
else:
gradient_mask_tensor = None
for epoch in range(done_epochs, cfg.max_epochs):
if engine.distributed and hasattr(train_data, 'train_sampler'):
train_data.train_sampler.set_epoch(epoch)
if epoch == done_epochs:
pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters))
else:
pbar = tqdm(range(iters_per_epoch))
if epoch == 0 and local_rank == 0:
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str='Init',
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
top1 = AvgMeter()
top5 = AvgMeter()
losses = AvgMeter()
discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs)
pbar.set_description('Train' + discrip_str)
for _ in pbar:
start_time = time.time()
data, label = load_cuda_data(train_data,
dataset_name=cfg.dataset_name)
# load_cuda_data(train_dataloader, cfg.dataset_name)
data_time = time.time() - start_time
if_accum_grad = ((iteration % cfg.grad_accum_iters) != 0)
train_net_time_start = time.time()
acc, acc5, loss = train_one_step(
model,
data,
label,
optimizer,
criterion,
if_accum_grad,
gradient_mask_tensor=gradient_mask_tensor,
lasso_keyword_to_strength=lasso_keyword_to_strength)
train_net_time_end = time.time()
if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters:
recorded_train_examples += cfg.global_batch_size
recorded_train_time += train_net_time_end - train_net_time_start
scheduler.step()
for module in model.modules():
if hasattr(module, 'set_cur_iter'):
module.set_cur_iter(iteration)
if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0:
for tag, value in zip(
tb_tags,
[acc.item(), acc5.item(),
loss.item()]):
tb_writer.add_scalars(tag, {'Train': value}, iteration)
top1.update(acc.item())
top5.update(acc5.item())
losses.update(loss.item())
if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS:
collected_train_loss_sum += loss.item()
collected_train_loss_count += 1
pbar_dic = OrderedDict()
pbar_dic['data-time'] = '{:.2f}'.format(data_time)
pbar_dic['cur_iter'] = iteration
pbar_dic['lr'] = scheduler.get_lr()[0]
pbar_dic['top1'] = '{:.5f}'.format(top1.mean)
pbar_dic['top5'] = '{:.5f}'.format(top5.mean)
pbar_dic['loss'] = '{:.5f}'.format(losses.mean)
pbar.set_postfix(pbar_dic)
iteration += 1
if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0:
engine.update_iteration(iteration)
if (not engine.distributed) or (engine.distributed and
engine.world_rank == 0):
engine.save_and_link_checkpoint(cfg.output_dir)
if iteration >= max_iters:
break
# do something after an epoch?
engine.update_iteration(iteration)
engine.save_latest_ckpt(cfg.output_dir)
if (epoch + 1) % save_hdf5_epochs == 0:
engine.save_hdf5(
os.path.join(cfg.output_dir,
'epoch-{}.hdf5'.format(epoch)))
if local_rank == 0 and \
cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0):
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str=discrip_str,
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
if iteration >= max_iters:
break
# do something after the training
if recorded_train_time > 0:
exp_per_sec = recorded_train_examples / recorded_train_time
else:
exp_per_sec = 0
engine.log(
'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}'
.format(int(TRAIN_SPEED_START * max_iters),
int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size,
recorded_train_examples, recorded_train_time, exp_per_sec))
if cfg.save_weights:
engine.save_checkpoint(cfg.save_weights)
print('NOTE: training finished, saved to {}'.format(
cfg.save_weights))
engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5'))
if collected_train_loss_count > 0:
engine.log(
'TRAIN LOSS collected over last {} epochs: {:.6f}'.format(
COLLECT_TRAIN_LOSS_EPOCHS,
collected_train_loss_sum / collected_train_loss_count))
def aofp_train_main(local_rank,
target_layers,
succ_strategy,
warmup_iterations,
aofp_batches_per_half,
flops_func,
cfg: BaseConfigByEpoch,
net=None,
train_dataloader=None,
val_dataloader=None,
show_variables=False,
convbuilder=None,
init_hdf5=None,
no_l2_keywords='depth',
gradient_mask=None,
use_nesterov=False,
tensorflow_style_init=False,
keyword_to_lr_mult=None,
auto_continue=False,
lasso_keyword_to_strength=None,
save_hdf5_epochs=10000,
remain_flops_ratio=0):
if no_l2_keywords is None:
no_l2_keywords = []
if type(no_l2_keywords) is not list:
no_l2_keywords = [no_l2_keywords]
ensure_dir(cfg.output_dir)
ensure_dir(cfg.tb_dir)
with Engine(local_rank=local_rank) as engine:
engine.setup_log(name='train',
log_dir=cfg.output_dir,
file_name='log.txt')
# ----------------------------- build model ------------------------------
if convbuilder is None:
convbuilder = ConvBuilder(base_config=cfg)
if net is None:
net_fn = get_model_fn(cfg.dataset_name, cfg.network_type)
model = net_fn(cfg, convbuilder)
else:
model = net
model = model.cuda()
# ----------------------------- model done ------------------------------
# ---------------------------- prepare data -------------------------
if train_dataloader is None:
train_data = create_dataset(cfg.dataset_name,
cfg.dataset_subset,
cfg.global_batch_size,
distributed=engine.distributed)
if cfg.val_epoch_period > 0 and val_dataloader is None:
val_data = create_dataset(cfg.dataset_name,
'val',
global_batch_size=100,
distributed=False)
engine.echo('NOTE: Data prepared')
engine.echo(
'NOTE: We have global_batch_size={} on {} GPUs, the allocated GPU memory is {}'
.format(cfg.global_batch_size, torch.cuda.device_count(),
torch.cuda.memory_allocated()))
# ----------------------------- data done --------------------------------
# ------------------------ parepare optimizer, scheduler, criterion -------
optimizer = get_optimizer(engine,
cfg,
model,
no_l2_keywords=no_l2_keywords,
use_nesterov=use_nesterov,
keyword_to_lr_mult=keyword_to_lr_mult)
scheduler = get_lr_scheduler(cfg, optimizer)
criterion = get_criterion(cfg).cuda()
# --------------------------------- done -------------------------------
engine.register_state(scheduler=scheduler,
model=model,
optimizer=optimizer)
if engine.distributed:
torch.cuda.set_device(local_rank)
engine.echo('Distributed training, device {}'.format(local_rank))
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],
broadcast_buffers=False,
)
else:
assert torch.cuda.device_count() == 1
engine.echo('Single GPU training')
if tensorflow_style_init:
init_as_tensorflow(model)
if cfg.init_weights:
engine.load_checkpoint(cfg.init_weights)
if init_hdf5:
engine.load_part('base_path.', init_hdf5)
if auto_continue:
assert cfg.init_weights is None
engine.load_checkpoint(get_last_checkpoint(cfg.output_dir))
if show_variables:
engine.show_variables()
# ------------ do training ---------------------------- #
engine.log("\n\nStart training with pytorch version {}".format(
torch.__version__))
iteration = engine.state.iteration
iters_per_epoch = num_iters_per_epoch(cfg)
max_iters = iters_per_epoch * cfg.max_epochs
tb_writer = SummaryWriter(cfg.tb_dir)
tb_tags = ['Top1-Acc', 'Top5-Acc', 'Loss']
model.train()
done_epochs = iteration // iters_per_epoch
last_epoch_done_iters = iteration % iters_per_epoch
if done_epochs == 0 and last_epoch_done_iters == 0:
engine.save_hdf5(os.path.join(cfg.output_dir, 'init.hdf5'))
recorded_train_time = 0
recorded_train_examples = 0
collected_train_loss_sum = 0
collected_train_loss_count = 0
if gradient_mask is not None:
gradient_mask_tensor = {}
for name, value in gradient_mask.items():
gradient_mask_tensor[name] = torch.Tensor(value).cuda()
else:
gradient_mask_tensor = None
######################### aofp
_init_interval = aofp_batches_per_half // len(target_layers)
layer_to_start_iter = {
i: (_init_interval * i + warmup_iterations)
for i in target_layers
}
print(
'the initial layer_to_start_iter = {}'.format(layer_to_start_iter))
# 0. get all the AOFPLayers
layer_idx_to_module = {}
for submodule in model.modules():
if hasattr(submodule, 'score_mask') or hasattr(
submodule, 't_value'):
layer_idx_to_module[submodule.conv_idx] = submodule
print(layer_idx_to_module)
######################################
for epoch in range(done_epochs, cfg.max_epochs):
if engine.distributed and hasattr(train_data, 'train_sampler'):
train_data.train_sampler.set_epoch(epoch)
if epoch == done_epochs:
pbar = tqdm(range(iters_per_epoch - last_epoch_done_iters))
else:
pbar = tqdm(range(iters_per_epoch))
if epoch == 0 and local_rank == 0:
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str='Init',
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
top1 = AvgMeter()
top5 = AvgMeter()
losses = AvgMeter()
discrip_str = 'Epoch-{}/{}'.format(epoch, cfg.max_epochs)
pbar.set_description('Train' + discrip_str)
for _ in pbar:
start_time = time.time()
data, label = load_cuda_data(train_data,
dataset_name=cfg.dataset_name)
# load_cuda_data(train_dataloader, cfg.dataset_name)
data_time = time.time() - start_time
if_accum_grad = ((iteration % cfg.grad_accum_iters) != 0)
train_net_time_start = time.time()
############ aofp
# 1. see if it is time to start on every layer
# 2. forward and accumulate
# 3. if a half on some layer is finished, do something
# ---- fetch its accumulated t vectors, analyze the first 'granu' elements
# ---- if good enough, set the base mask, reset the search space
# ---- elif granu == 1, do nothing
# ---- else, granu /= 2, reset the search space
for layer_idx, start_iter in layer_to_start_iter.items():
if start_iter == iteration:
layer_idx_to_module[layer_idx].start_aofp(iteration)
acc, acc5, loss = train_one_step(
model,
data,
label,
optimizer,
criterion,
if_accum_grad,
gradient_mask_tensor=gradient_mask_tensor,
lasso_keyword_to_strength=lasso_keyword_to_strength)
for layer_idx, aofp_layer in layer_idx_to_module.items():
# accumulate
if layer_idx not in succ_strategy:
continue
follow_layer_idx = succ_strategy[layer_idx]
if follow_layer_idx not in layer_idx_to_module:
continue
t_value = layer_idx_to_module[follow_layer_idx].t_value
aofp_layer.accumulate_t_value(t_value)
if aofp_layer.finished_a_half(iteration):
aofp_layer.halve_or_stop(iteration)
###################################
train_net_time_end = time.time()
if iteration > TRAIN_SPEED_START * max_iters and iteration < TRAIN_SPEED_END * max_iters:
recorded_train_examples += cfg.global_batch_size
recorded_train_time += train_net_time_end - train_net_time_start
scheduler.step()
for module in model.modules():
if hasattr(module, 'set_cur_iter'):
module.set_cur_iter(iteration)
if iteration % cfg.tb_iter_period == 0 and engine.world_rank == 0:
for tag, value in zip(
tb_tags,
[acc.item(), acc5.item(),
loss.item()]):
tb_writer.add_scalars(tag, {'Train': value}, iteration)
top1.update(acc.item())
top5.update(acc5.item())
losses.update(loss.item())
if epoch >= cfg.max_epochs - COLLECT_TRAIN_LOSS_EPOCHS:
collected_train_loss_sum += loss.item()
collected_train_loss_count += 1
pbar_dic = OrderedDict()
pbar_dic['data-time'] = '{:.2f}'.format(data_time)
pbar_dic['cur_iter'] = iteration
pbar_dic['lr'] = scheduler.get_lr()[0]
pbar_dic['top1'] = '{:.5f}'.format(top1.mean)
pbar_dic['top5'] = '{:.5f}'.format(top5.mean)
pbar_dic['loss'] = '{:.5f}'.format(losses.mean)
pbar.set_postfix(pbar_dic)
iteration += 1
if iteration >= max_iters or iteration % cfg.ckpt_iter_period == 0:
engine.update_iteration(iteration)
if (not engine.distributed) or (engine.distributed and
engine.world_rank == 0):
engine.save_and_link_checkpoint(cfg.output_dir)
if iteration >= max_iters:
break
# do something after an epoch?
engine.update_iteration(iteration)
engine.save_latest_ckpt(cfg.output_dir)
if (epoch + 1) % save_hdf5_epochs == 0:
engine.save_hdf5(
os.path.join(cfg.output_dir,
'epoch-{}.hdf5'.format(epoch)))
if local_rank == 0 and \
cfg.val_epoch_period > 0 and (epoch >= cfg.max_epochs - 10 or epoch % cfg.val_epoch_period == 0):
val_during_train(epoch=epoch,
iteration=iteration,
tb_tags=tb_tags,
engine=engine,
model=model,
val_data=val_data,
criterion=criterion,
descrip_str=discrip_str,
dataset_name=cfg.dataset_name,
test_batch_size=TEST_BATCH_SIZE,
tb_writer=tb_writer)
cur_deps = np.array(cfg.deps)
for submodule in model.modules():
if hasattr(submodule, 'base_mask'):
cur_deps[submodule.conv_idx] = np.sum(
submodule.base_mask.cpu().numpy() == 1)
origin_flops = flops_func(cfg.deps)
cur_flops = flops_func(cur_deps)
remain_ratio = cur_flops / origin_flops
if local_rank == 0:
print('##########################')
print('origin deps ', cfg.deps)
print('cur deps ', cur_deps)
print('remain flops ratio = ', remain_ratio, 'the target is ',
remain_flops_ratio)
print('##########################')
if remain_ratio < remain_flops_ratio:
break
if iteration >= max_iters:
break
# do something after the training
if recorded_train_time > 0:
exp_per_sec = recorded_train_examples / recorded_train_time
else:
exp_per_sec = 0
engine.log(
'TRAIN speed: from {} to {} iterations, batch_size={}, examples={}, total_net_time={:.4f}, examples/sec={}'
.format(int(TRAIN_SPEED_START * max_iters),
int(TRAIN_SPEED_END * max_iters), cfg.global_batch_size,
recorded_train_examples, recorded_train_time, exp_per_sec))
if cfg.save_weights:
engine.save_checkpoint(cfg.save_weights)
print('NOTE: training finished, saved to {}'.format(
cfg.save_weights))
engine.save_hdf5(os.path.join(cfg.output_dir, 'finish.hdf5'))
if collected_train_loss_count > 0:
engine.log(
'TRAIN LOSS collected over last {} epochs: {:.6f}'.format(
COLLECT_TRAIN_LOSS_EPOCHS,
collected_train_loss_sum / collected_train_loss_count))
final_deps = aofp_prune(model,
origin_deps=cfg.deps,
succ_strategy=succ_strategy,
save_path=os.path.join(cfg.output_dir,
'finish_pruned.hdf5'))
origin_flops = flops_func(cfg.deps)
cur_flops = flops_func(final_deps)
engine.log(
'##################################################################'
)
engine.log(cfg.network_type)
engine.log('origin width: {} , flops {} '.format(
cfg.deps, origin_flops))
engine.log('final width: {}, flops {} '.format(final_deps, cur_flops))
engine.log('flops reduction: {}'.format(1 - cur_flops / origin_flops))
return final_deps