def init_dataloader(config):
trainloader, testloader = get_dataloader(
dataset=config.dataset,
train_batch_size=config.batch_size,
test_batch_size=256,
returnset=config.data_distributed)
return trainloader, testloader
def main(config, args):
# init logger
classes = {
'cifar10': 10,
'cifar100': 100,
'mnist': 10,
'tiny_imagenet': 200
}
logger, writer = init_logger(config, args)
best_acc_vec = []
test_acc_vec_vec = []
for n_runs in range(1):
if args.sigma_w2 != None and n_runs != 0:
break
# build model
model = get_network(config.network,
config.depth,
config.dataset,
use_bn=config.get('use_bn', args.bn),
scaled=args.scaled_init,
act=args.act)
mask = None
mb = ModelBase(config.network, config.depth, config.dataset, model)
mb.cuda()
if mask is not None:
mb.register_mask(mask)
ratio_vec_ = print_mask_information(mb, logger)
# preprocessing
# ====================================== get dataloader ======================================
trainloader, testloader = get_dataloader(config.dataset,
config.batch_size, 256, 4)
# ====================================== fetch configs ======================================
ckpt_path = config.checkpoint_dir
num_iterations = config.iterations
if args.target_ratio == None:
target_ratio = config.target_ratio
else:
target_ratio = args.target_ratio
normalize = config.normalize
# ====================================== fetch exception ======================================
exception = get_exception_layers(
mb.model, str_to_list(config.exception, ',', int))
logger.info('Exception: ')
for idx, m in enumerate(exception):
logger.info(' (%d) %s' % (idx, m))
# ====================================== fetch training schemes ======================================
ratio = 1 - (1 - target_ratio)**(1.0 / num_iterations)
learning_rates = str_to_list(config.learning_rate, ',', float)
weight_decays = str_to_list(config.weight_decay, ',', float)
training_epochs = str_to_list(config.epoch, ',', int)
logger.info(
'Normalize: %s, Total iteration: %d, Target ratio: %.2f, Iter ratio %.4f.'
% (normalize, num_iterations, target_ratio, ratio))
logger.info('Basic Settings: ')
for idx in range(len(learning_rates)):
logger.info(' %d: LR: %.5f, WD: %.5f, Epochs: %d' %
(idx, learning_rates[idx], weight_decays[idx],
training_epochs[idx]))
# ====================================== start pruning ======================================
iteration = 0
for _ in range(1):
logger.info(
'** Target ratio: %.4f, iter ratio: %.4f, iteration: %d/%d.' %
(target_ratio, ratio, 1, num_iterations))
# mb.model.apply(weights_init)
print('#' * 40)
print('USING {} INIT SCHEME'.format(args.init))
print('#' * 40)
if args.init == 'kaiming_xavier':
mb.model.apply(weights_init_kaiming_xavier)
elif args.init == 'kaiming':
if args.act == 'relu' or args.act == 'elu':
mb.model.apply(weights_init_kaiming_relu)
elif args.act == 'tanh':
mb.model.apply(weights_init_kaiming_tanh)
elif args.init == 'xavier':
mb.model.apply(weights_init_xavier)
elif args.init == 'EOC':
mb.model.apply(weights_init_EOC)
elif args.init == 'ordered':
def weights_init_ord(m):
if isinstance(m, nn.Conv2d):
ord_weights(m.weight, sigma_w2=args.sigma_w2)
if m.bias is not None:
ord_bias(m.bias)
elif isinstance(m, nn.Linear):
ord_weights(m.weight, sigma_w2=args.sigma_w2)
if m.bias is not None:
ord_bias(m.bias)
elif isinstance(m, nn.BatchNorm2d):
# Note that BN's running_var/mean are
# already initialized to 1 and 0 respectively.
if m.weight is not None:
m.weight.data.fill_(1.0)
if m.bias is not None:
m.bias.data.zero_()
mb.model.apply(weights_init_ord)
else:
raise NotImplementedError
print("=> Applying weight initialization(%s)." %
config.get('init_method', 'kaiming'))
print("Iteration of: %d/%d" % (iteration, num_iterations))
if config.pruner == 'SNIP':
print('=> Using SNIP')
masks, scaled_masks = SNIP(
mb.model,
ratio,
trainloader,
'cuda',
num_classes=classes[config.dataset],
samples_per_class=config.samples_per_class,
num_iters=config.get('num_iters', 1),
scaled_init=args.scaled_init)
elif config.pruner == 'GraSP':
print('=> Using GraSP')
masks, scaled_masks = GraSP(
mb.model,
ratio,
trainloader,
'cuda',
num_classes=classes[config.dataset],
samples_per_class=config.samples_per_class,
num_iters=config.get('num_iters', 1),
scaled_init=args.scaled_init)
iteration = 0
################################################################################
_masks = None
_masks_scaled = None
if not args.bn:
# build model that has the same weights as the pruned network but with BN now !
model2 = get_network(config.network,
config.depth,
config.dataset,
use_bn=config.get('use_bn', True),
scaled=args.scaled_init,
act=args.act)
weights_temp = []
for layer_old in mb.model.modules():
if isinstance(layer_old, nn.Conv2d) or isinstance(
layer_old, nn.Linear):
weights_temp.append(layer_old.weight)
idx = 0
for layer_new in model2.modules():
if isinstance(layer_new, nn.Conv2d) or isinstance(
layer_new, nn.Linear):
layer_new.weight.data = weights_temp[idx]
idx += 1
# Creating a base model with BN included now
mb = ModelBase(config.network, config.depth, config.dataset,
model2)
mb.cuda()
_masks = dict()
_masks_scaled = dict()
layer_keys_new = []
for layer in (mb.model.modules()):
if isinstance(layer, nn.Conv2d) or isinstance(
layer, nn.Linear):
layer_keys_new.append(layer)
for new_keys, old_keys in zip(layer_keys_new, masks.keys()):
_masks[new_keys] = masks[old_keys]
if args.scaled_init:
_masks_scaled[new_keys] = scaled_masks[old_keys]
################################################################################
if _masks == None:
_masks = masks
_masks_scaled = scaled_masks
# ========== register mask ==================
mb.register_mask(_masks)
## ========== debugging ==================
if args.scaled_init:
if config.network == 'vgg':
print('scaling VGG')
mb.scaling_weights(_masks_scaled)
# ========== save pruned network ============
logger.info('Saving..')
state = {
'net': mb.model,
'acc': -1,
'epoch': -1,
'args': config,
'mask': mb.masks,
'ratio': mb.get_ratio_at_each_layer()
}
path = os.path.join(
ckpt_path, 'prune_%s_%s%s_r%s_it%d.pth.tar' %
(config.dataset, config.network, config.depth, target_ratio,
iteration))
torch.save(state, path)
# ========== print pruning details ============
logger.info('**[%d] Mask and training setting: ' % iteration)
ratio_vec_ = print_mask_information(mb, logger)
logger.info(' LR: %.5f, WD: %.5f, Epochs: %d' %
(learning_rates[iteration], weight_decays[iteration],
training_epochs[iteration]))
results_path = config.summary_dir + args.init + '_sp' + str(
args.target_ratio).replace('.', '_')
if args.scaled_init:
results_path += '_scaled'
if args.bn:
results_path += '_bn'
if args.sigma_w2 != None and args.init == 'ordered':
results_path += '_sgw2{}'.format(args.sigma_w2).replace(
'.', '_')
results_path += '_' + args.act + '_' + str(config.depth)
print('saving the ratios')
print(results_path)
if not os.path.isdir(results_path): os.mkdir(results_path)
np.save(results_path + '/ratios_pruned{}'.format(args.seed_tiny),
np.array(ratio_vec_))
# if args.sigma_w2 != None:
# break
# ========== finetuning =======================
best_acc, test_acc_vec = train_once(
mb=mb,
net=mb.model,
trainloader=trainloader,
testloader=testloader,
writer=writer,
config=config,
ckpt_path=ckpt_path,
learning_rate=learning_rates[iteration],
weight_decay=weight_decays[iteration],
num_epochs=training_epochs[iteration],
iteration=iteration,
logger=logger,
args=args)
best_acc_vec.append(best_acc)
test_acc_vec_vec.append(test_acc_vec)
np.save(results_path + '/best_acc{}'.format(args.seed_tiny),
np.array(best_acc_vec))
np.save(results_path + '/test_acc{}'.format(args.seed_tiny),
np.array(test_acc_vec_vec))
}
act = torch.nn.Sigmoid() if args.activation is None else act_dict[
args.activation]
# init model
encoder_sizes = [28 * 28, 1000, 500, 250, 30]
decoder_sizes = [30, 250, 500, 1000, 28 * 28]
net = deep_autoencoder(encoder_sizes=encoder_sizes,
decoder_sizes=decoder_sizes,
activation=act).to(args.device)
# init dataloader
trainloader, testloader = get_dataloader(dataset=args.dataset,
train_batch_size=args.batch_size,
test_batch_size=256)
# init optimizer
optim_name = args.optimizer.lower()
tag = optim_name
optimizer = get_optimizer(optim_name, net, args)
# init lr scheduler
lr_scheduler = get_lr_scheduler(optimizer, args)
# init criterion
criterion = torch.nn.BCEWithLogitsLoss()
# init summary writter
log_dir = get_log_dir(optim_name, args)
if not os.path.isdir(log_dir):
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument("--config",
default=None,
type=str,
required=True,
help="the training config file")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument("--multi_task",
action="store_true",
help="training with multi task schema")
parser.add_argument("--debug",
action="store_true",
help="in debug mode, will not enable wandb log")
parser.add_argument("--use_wandb",
action="store_true",
help="whether or not use wandb")
args = parser.parse_args()
cfg = parse_cfg(pathlib.Path(args.config))
# set CUDA_VISIBLE_DEVICES and get num_gpus
if args.local_rank == -1: # not distributed
os.environ["CUDA_VISIBLE_DEVICES"] = cfg["system"][
"cuda_visible_devices"]
num_gpus = torch.cuda.device_count()
args.distributed = False
else: # distributed
torch.cuda.set_device(args.local_rank)
num_gpus = 1
args.distributed = True
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
logger.info(
"num_gpus: {}, distributed training: {}, 16-bits training: {}".format(
num_gpus, bool(args.local_rank != -1), cfg["train"]["fp16"]))
cudnn.benchmark = True
cfg["train"]["output_dir"] = cfg["train"]["output_dir"] + "/" + \
cfg["train"]["task_name"] + "_" + \
cfg["train"]["model_name"] + "_" + \
cfg["data"]["corpus"]
output_dir_pl = pathlib.Path(cfg["train"]["output_dir"])
if output_dir_pl.exists():
logger.warn(
"output directory ({}) already exists, continue after 2 seconds..."
.format(output_dir_pl))
time.sleep(2)
else:
output_dir_pl.mkdir(parents=True, exist_ok=True)
if not args.debug and args.use_wandb:
config_dictionary = dict(yaml=cfg, params=args)
wandb.init(config=config_dictionary,
project="nlp-task",
dir=cfg["train"]["output_dir"])
wandb.run.name = cfg["data"]["corpus"] + '-' + cfg["train"][
"pretrained_tag"] + '-' + time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
wandb.config.update(args)
wandb.run.save()
if cfg["optimizer"]["gradient_accumulation_steps"] < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(cfg["optimizer"]["gradient_accumulation_steps"]))
# true batch_size in training
cfg["train"]["batch_size"] = cfg["train"]["batch_size"] // cfg[
"optimizer"]["gradient_accumulation_steps"]
# the type of label_map is bidict
# label_map[x] = xx, label_map.inv[xx] = x
label_map, num_labels = get_label_map(cfg)
tokenizer, model = get_tokenizer_and_model(cfg, label_map, num_labels)
# check model details on wandb
if not args.debug and args.use_wandb:
wandb.watch(model)
num_examples, train_dataloader = get_dataloader(cfg,
tokenizer,
num_labels,
"train",
debug=args.debug)
_, eval_dataloader = get_dataloader(cfg,
tokenizer,
num_labels,
"dev",
debug=args.debug)
# total training steps (including multi epochs)
num_training_steps = int(
len(train_dataloader) //
cfg["optimizer"]["gradient_accumulation_steps"] *
cfg["train"]["train_epochs"])
optimizer = AdamW(params=model.parameters(), lr=cfg["optimizer"]["lr"])
lr_scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=cfg["optimizer"]["num_warmup_steps"],
num_training_steps=num_training_steps)
scaler = None
model = model.cuda()
if cfg["train"]["fp16"] and _use_apex:
logger.error("using apex amp for fp16...")
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
elif cfg["train"]["fp16"] and _use_native_amp:
logger.error("using pytorch native amp for fp16...")
scaler = torch.cuda.amp.GradScaler()
elif cfg["train"]["fp16"] and (_use_apex is False
and _use_native_amp is False):
logger.error("your environment DO NOT support fp16 training...")
exit()
if cfg["system"]["distributed"]:
# TODO distributed debug
model.cuda(args.local_rank)
from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model, device_ids=[args.local_rank])
elif num_gpus > 1:
model = torch.nn.DataParallel(model)
# Train
logger.info("start training on train set")
epoch = 0
best_score = -1
for _ in trange(int(cfg["train"]["train_epochs"]), desc="Epoch"):
best = False
# train loop in one epoch
train_loop(cfg, model, train_dataloader, optimizer, lr_scheduler,
num_gpus, epoch, scaler, args.debug, args.use_wandb)
# begin to evaluate
logger.info("running evaluation on dev set")
score = eval_loop(cfg, tokenizer, model, eval_dataloader, label_map,
args.debug, args.use_wandb)
if best_score < score:
best_score = score
best = True
# Save a trained model and the associated configuration
save_model(cfg, tokenizer, model, best)
epoch += 1
# Test Eval
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
logger.info("running evaluation on final test set")
# TODO add stand alone test set
_, eval_dataloader = get_dataloader(cfg,
tokenizer,
num_labels,
"dev",
debug=args.debug)
score = eval_loop(cfg, tokenizer, model, eval_dataloader, label_map,
args.debug, args.use_wandb)
def init_dataloader(config):
trainloader, testloader = get_dataloader(dataset=config.dataset,
train_batch_size=config.batch_size,
test_batch_size=128)
return trainloader, testloader
def main(config):
# init logger
classes = {
'cifar10': 10,
'cifar100': 100,
'mnist': 10,
'tiny_imagenet': 200
}
logger, writer = init_logger(config)
# build model
model = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', True))
mask = None
mb = ModelBase(config.network, config.depth, config.dataset, model)
mb.cuda()
if mask is not None:
mb.register_mask(mask)
print_mask_information(mb, logger)
# preprocessing
# ====================================== get dataloader ======================================
trainloader, testloader = get_dataloader(config.dataset, config.batch_size, 256, 4, root='/home/wzn/PycharmProjects/GraSP/data')
# ====================================== fetch configs ======================================
ckpt_path = config.checkpoint_dir
num_iterations = config.iterations
target_ratio = config.target_ratio
normalize = config.normalize
# ====================================== fetch exception ======================================
# exception = get_exception_layers(mb.model, str_to_list(config.exception, ',', int))
# logger.info('Exception: ')
#
# for idx, m in enumerate(exception):
# logger.info(' (%d) %s' % (idx, m))
# ====================================== fetch training schemes ======================================
ratio = 1 - (1 - target_ratio) ** (1.0 / num_iterations)
learning_rates = str_to_list(config.learning_rate, ',', float)
weight_decays = str_to_list(config.weight_decay, ',', float)
training_epochs = str_to_list(config.epoch, ',', int)
logger.info('Normalize: %s, Total iteration: %d, Target ratio: %.2f, Iter ratio %.4f.' %
(normalize, num_iterations, target_ratio, ratio))
logger.info('Basic Settings: ')
for idx in range(len(learning_rates)):
logger.info(' %d: LR: %.5f, WD: %.5f, Epochs: %d' % (idx,
learning_rates[idx],
weight_decays[idx],
training_epochs[idx]))
# ====================================== start pruning ======================================
iteration = 0
for _ in range(1):
# logger.info('** Target ratio: %.4f, iter ratio: %.4f, iteration: %d/%d.' % (target_ratio,
# ratio,
# 1,
# num_iterations))
mb.model.apply(weights_init)
print("=> Applying weight initialization(%s)." % config.get('init_method', 'kaiming'))
# print("Iteration of: %d/%d" % (iteration, num_iterations))
# masks = GraSP(mb.model, ratio, trainloader, 'cuda',
# num_classes=classes[config.dataset],
# samples_per_class=config.samples_per_class,
# num_iters=config.get('num_iters', 1))
# iteration = 0
# print('=> Using GraSP')
# # ========== register mask ==================
# mb.register_mask(masks)
# # ========== save pruned network ============
# logger.info('Saving..')
# state = {
# 'net': mb.model,
# 'acc': -1,
# 'epoch': -1,
# 'args': config,
# 'mask': mb.masks,
# 'ratio': mb.get_ratio_at_each_layer()
# }
# path = os.path.join(ckpt_path, 'prune_%s_%s%s_r%s_it%d.pth.tar' % (config.dataset,
# config.network,
# config.depth,
# config.target_ratio,
# iteration))
# torch.save(state, path)
# # ========== print pruning details ============
# logger.info('**[%d] Mask and training setting: ' % iteration)
# print_mask_information(mb, logger)
# logger.info(' LR: %.5f, WD: %.5f, Epochs: %d' %
# (learning_rates[iteration], weight_decays[iteration], training_epochs[iteration]))
# ========== finetuning =======================
train_once(mb=mb,
net=mb.model,
trainloader=trainloader,
testloader=testloader,
writer=writer,
config=config,
ckpt_path=ckpt_path,
learning_rate=learning_rates[iteration],
weight_decay=weight_decays[iteration],
num_epochs=training_epochs[iteration],
iteration=iteration,
logger=logger)
