def configure(model, compress):
model = model.cuda()
compression_scheduler = distiller.CompressionScheduler(model)
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
if compress:
source = compress
compression_scheduler = distiller.CompressionScheduler(model)
distiller.config.file_config(model, optimizer, compress, compression_scheduler,)
parser = argparse.ArgumentParser()
distiller.knowledge_distillation.add_distillation_args(parser)
CONFIG_FILE = '.config_ipynb'
if os.path.isfile(CONFIG_FILE):
with open(CONFIG_FILE) as f:
sys.argv = f.read().split()
else:
sys.argv = ['resnet.py','--kd-resume', 'net56_cifar.pth.tar','--kd-teacher',None,
'--kd-start-epoch', 0,
'--kd-student-wt','0.5',
'--kd-teacher-wt', '0.0',
'--kd-distill-wt', '0.5',
]
args = parser.parse_args()
args.kd_policy = None
epochs = 30
if args.kd_teacher:
if args.kd_resume:
teacher = torch.load(args.kd_resume)
# Create policy and add to scheduler
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=epochs,
frequency=1)
return model, compression_scheduler, epochs, optimizer, criterion, args
def init_knowledge_distillation(args, model, compression_scheduler):
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained, args.dataset, args.kd_teacher, device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if args.epochs is None:
args.epochs = 90
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(args.compress, msglogger.logdir, gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
ending_epoch = args.epochs
perf_scores_history = []
if args.evaluate:
args.deterministic = True
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
distiller.set_deterministic() # Use a well-known seed, for repeatability of experiments
else:
# Turn on CUDNN benchmark mode for best performance. This is usually "safe" for image
# classification models, as the input sizes don't change during the run
# See here: https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
cudnn.benchmark = True
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
raise ValueError('ERROR: Argument --gpus must be a comma-separated list of integers only')
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
raise ValueError('ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained, args.dataset, args.arch,
parallel=not args.load_serialized, device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s', args.earlyexit_thresholds)
# TODO(barrh): args.deprecated_resume is deprecated since v0.3.1
if args.deprecated_resume:
msglogger.warning('The "--resume" flag is deprecated. Please use "--resume-from=YOUR_PATH" instead.')
if not args.reset_optimizer:
msglogger.warning('If you wish to also reset the optimizer, call with: --reset-optimizer')
args.reset_optimizer = True
args.resumed_checkpoint_path = args.deprecated_resume
# We can optionally resume from a checkpoint
optimizer = None
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model, args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info('\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0')
# Define loss function (criterion)
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is None:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer, pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
activations_collectors = create_activation_stats_collectors(model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info('\tStats will be collected for {:.1%} of test dataset'.format(args.qe_calibration))
msglogger.info('\tSetting constant seeds and converting model to serialized execution')
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(create_quantization_stats_collector(model))
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size, args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0], args.sensitivity_range[1], args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger, args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger, activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress, compression_scheduler,
(start_epoch-1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, \
"You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model, compression_scheduler.zeros_mask_dict, args.arch, args.dataset, optimizer=None)
apputils.save_checkpoint(0, args.arch, model, optimizer=None, scheduler=compression_scheduler,
name="{}_thinned".format(args.resumed_checkpoint_path.replace(".pth.tar", "")),
dir=msglogger.logdir)
print("Note: your model may have collapsed to random inference, so you may want to fine-tune")
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained, args.dataset, args.kd_teacher, device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'.format(
start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
for epoch in range(start_epoch, ending_epoch):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch,
metrics=(vloss if (epoch != start_epoch) else 10**6))
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], args=args)
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(epoch, "train", loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(distiller.masks_sparsity_tbl_summary(model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion, [pylogger], args, epoch)
distiller.log_activation_statsitics(epoch, "valid", loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss),
('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats, None, epoch, steps_completed=0, total_steps=1, log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint
update_training_scores_history(perf_scores_history, model, top1, top5, epoch, args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {'current_top1': top1,
'best_top1': perf_scores_history[0].top1,
'best_epoch': perf_scores_history[0].epoch}
apputils.save_checkpoint(epoch, args.arch, model, optimizer=optimizer, scheduler=compression_scheduler,
extras=checkpoint_extras, is_best=is_best, name=args.name, dir=msglogger.logdir)
# Finally run results on the test set
test(test_loader, model, criterion, [pylogger], activations_collectors, args=args)
def main():
global msglogger
check_pytorch_version()
args = parser.parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(sys.argv, gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_epochs = [distiller.MutableNamedTuple({'epoch': 0, 'top1': 0, 'sparsity': 0})
for i in range(args.num_best_scores)]
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error('ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1')
exit(1)
# Use a well-known seed, for repeatability of experiments
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
cudnn.deterministic = True
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error('ERROR: Argument --gpus must be a comma-separated list of integers only')
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error('ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
if 'cinic' in args.arch:
args.dataset = 'cinic10'
else:
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset in ['cifar10', 'cinic10'] else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
#model = create_model(args.pretrained, args.dataset, args.arch,
# parallel=not args.load_serialized, device_ids=args.gpus)
model = create_model(False, args.dataset, args.arch, device_ids=args.gpus) # Get arch state_dict
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s', args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
if args.resume:
#model, compression_scheduler, start_epoch = apputils.load_checkpoint(
# model, chkpt_file=args.resume)
# Load Pre-trained Model
chkpt_file=args.resume
print("=> loading checkpoint %s" % chkpt_file)
checkpoint = torch.load(chkpt_file)
model.load_state_dict(checkpoint['net'])
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.ADC:
return automated_deep_compression(model, criterion, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.workers, args.validation_size, args.deterministic)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler), len(test_loader.sampler))
activations_collectors = create_activation_stats_collectors(model, collection_phase=args.activation_stats)
if args.sensitivity is not None:
return sensitivity_analysis(model, criterion, test_loader, pylogger, args)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger, activations_collectors, args)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(model, optimizer, args.compress)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.cuda()
else:
compression_scheduler = distiller.CompressionScheduler(model)
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained, args.dataset, args.kd_teacher, device_ids=args.gpus)
if args.kd_resume:
teacher, _, _ = apputils.load_checkpoint(teacher, chkpt_file=args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader, model, criterion, optimizer, epoch, compression_scheduler,
loggers=[tflogger, pylogger], args=args)
distiller.log_weights_sparsity(model, epoch, loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(epoch, "train", loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(distiller.masks_sparsity_tbl_summary(model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion, [pylogger], args, epoch)
distiller.log_activation_statsitics(epoch, "valid", loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss),
('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats, None, epoch, steps_completed=0, total_steps=1, log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# remember best top1 and save checkpoint
#sparsity = distiller.model_sparsity(model)
is_best = top1 > best_epochs[0].top1
if is_best:
best_epochs[0].epoch = epoch
best_epochs[0].top1 = top1
#best_epoch.sparsity = sparsity
best_epochs = sorted(best_epochs, key=lambda score: score.top1)
for score in reversed(best_epochs):
if score.top1 > 0:
msglogger.info('==> Best Top1: %.3f on Epoch: %d', score.top1, score.epoch)
apputils.save_checkpoint(epoch, args.arch, model, optimizer, compression_scheduler,
best_epochs[0].top1, is_best, args.name, msglogger.logdir)
# Finally run results on the test set
test(test_loader, model, criterion, [pylogger], activations_collectors, args=args)
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument("--data_dir",
default=None,
type=str,
required=True,
help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
parser.add_argument("--bert_model", default=None, type=str, required=True,
choices=[
"bert-base-uncased",
"bert-large-uncased",
"bert-base-cased",
"bert-large-cased",
"bert-base-multilingual-uncased",
"bert-base-multilingual-cased",
"bert-base-chinese",
],
help="Bert pre-trained model selected in the list")
parser.add_argument("--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train.")
parser.add_argument("--output_dir",
default=None,
type=str,
required=True,
help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument("--labels",
nargs='+',
default=['0', '1'],
help="labels")
## Other parameters
parser.add_argument("--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--do_test",
action='store_true',
help="Whether to run eval on the test set.")
parser.add_argument("--do_distill",
action='store_true',
help="Whether to run distillation.")
parser.add_argument("--blendcnn_channels",
nargs='+',
default=(100,) * 8,
help="BlendCNN channels.")
parser.add_argument("--blendcnn_act",
default='relu',
choices=list(ACT2FN.keys()),
help="BlendCNN activation function.")
parser.add_argument('--blendcnn_dropout',
action='store_true',
help="Whether to use dropout in BlendCNN")
parser.add_argument('--blendcnn_pair',
action='store_true',
help="Whether to use BlendCNNForSequencePairClassification")
parser.add_argument("--export_onnx",
action='store_true',
help="Whether to export model to onnx format.")
parser.add_argument("--onnx_framework",
choices=[
"caffe2",
],
help="Select the ONNX framework to run eval")
parser.add_argument("--eval_interval",
default=1000,
type=int,
help="Specify eval interval during training.")
parser.add_argument("--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument('--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument('--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument('--loss_scale',
type=float, default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
distiller.knowledge_distillation.add_distillation_args(parser)
args = parser.parse_args()
processors = {
"cola": ColaProcessor,
"mnli": MnliProcessor,
"mrpc": MrpcProcessor,
"custom": lambda: CustomProcessor(args.labels),
}
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
args.gradient_accumulation_steps))
args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not any((args.do_train, args.do_eval, args.do_test, args.do_distill, args.export_onnx)):
raise ValueError("At least one of `do_train`, `do_eval`, `do_test`, `do_distill`, `export_onnx` must be True.")
if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train:
raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
os.makedirs(args.output_dir, exist_ok=True)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels()
num_labels = len(label_list)
tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
global_step = 0
loss = 0
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
onnx_model_file = os.path.join(args.output_dir, "model.onnx")
eval_data = None
if args.do_train:
model = BertForSequenceClassification.from_pretrained(args.bert_model,
cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(
args.local_rank),
num_labels=num_labels)
model = convert_model(args, model, device, n_gpu)
tensorboard_log_dir = os.path.join(args.output_dir, './log')
os.makedirs(tensorboard_log_dir, exist_ok=True)
tensorboard_logger = SummaryWriter(tensorboard_log_dir)
if args.do_eval and do_eval_or_test(args) and eval_data is None:
eval_data = prepare(args, processor, label_list, tokenizer, 'dev')
global_step, loss = train(args,
model,
output_model_file,
processor,
label_list,
tokenizer,
device,
n_gpu,
tensorboard_logger,
eval_data)
model_config = None
model_embeddings = None
if args.onnx_framework is None:
# Load a trained model that you have fine-tuned
if os.path.exists(output_model_file):
model_state_dict = torch.load(output_model_file, map_location=lambda storage, loc: storage)
else:
model_state_dict = None
model = BertForSequenceClassification.from_pretrained(args.bert_model,
state_dict=model_state_dict,
num_labels=num_labels)
model_config = copy.deepcopy(model.config)
model_embeddings = model.bert.embeddings
model = convert_model(args, model, device, n_gpu)
else:
import onnx
model = onnx.load(onnx_model_file)
onnx.checker.check_model(model)
if args.do_distill:
assert model_config is not None
assert model_embeddings is not None
output_distilled_model_file = os.path.join(args.output_dir, DISTILLER_WEIGHTS_NAME)
teacher = model
model_config.hidden_act = args.blendcnn_act
if args.blendcnn_pair:
student = BlendCNNForSequencePairClassification(model_config,
num_labels=num_labels,
channels=(model_config.hidden_size,) +
args.blendcnn_channels,
n_hidden_dense=(model_config.hidden_size,) * 2,
use_dropout=args.blendcnn_dropout)
else:
student = BlendCNN(model_config,
num_labels=num_labels,
channels=(model_config.hidden_size,) + args.blendcnn_channels,
n_hidden_dense=(model_config.hidden_size,) * 2,
use_dropout=args.blendcnn_dropout)
student.embeddings.load_state_dict(model_embeddings.state_dict())
student = convert_model(args, student, device, 1)
if os.path.exists(output_distilled_model_file):
logger.info(
'Loading existing distilled model {}, skipping distillation'.format(output_distilled_model_file))
student.load_state_dict(torch.load(output_distilled_model_file))
else:
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(student, teacher, args.kd_temp, dlw)
tensorboard_log_dir = os.path.join(args.output_dir, './log')
os.makedirs(tensorboard_log_dir, exist_ok=True)
tensorboard_logger = SummaryWriter(tensorboard_log_dir)
if args.do_eval and do_eval_or_test(args) and eval_data is None:
eval_data = prepare(args, processor, label_list, tokenizer, 'dev')
global_step, loss = distill(args,
output_distilled_model_file,
processor,
label_list,
tokenizer,
device,
n_gpu,
tensorboard_logger,
eval_data)
model = student
if do_eval_or_test(args):
result = {
'global_step': global_step,
'loss': loss
}
model.float()
name = '_distiller' if args.do_distill else ''
if args.do_eval:
if eval_data is None:
eval_data = prepare(args, processor, label_list, tokenizer, 'dev')
eval_loss, eval_accuracy, eval_probs = eval(args, model, eval_data, device, verbose=True)
np.savetxt(os.path.join(args.output_dir, 'dev{}_probs.npy'.format(name)), eval_probs)
result.update({
'dev{}_loss'.format(name): eval_loss,
'dev{}_accuracy'.format(name): eval_accuracy,
})
if args.do_test:
eval_data = prepare(args, processor, label_list, tokenizer, 'test')
eval_loss, eval_accuracy, eval_probs = eval(args, model, eval_data, device, verbose=True)
np.savetxt(os.path.join(args.output_dir, 'test{}_probs.npy'.format(name)), eval_probs)
result.update({
'test{}_loss'.format(name): eval_loss,
'test{}_accuracy'.format(name): eval_accuracy,
})
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
if args.export_onnx:
if not env_enabled(ENV_OPENAIGPT_GELU) or not env_enabled(ENV_DISABLE_APEX):
raise ValueError('Both {} and {} must be 1 to properly export ONNX.'.format(ENV_OPENAIGPT_GELU,
ENV_DISABLE_APEX))
if not isinstance(model, torch.nn.Module):
raise ValueError('model is not an instance of torch.nn.Module.')
import onnx
import onnx.utils
import onnx.optimizer
dummy_input = get_dummy_input(args, processor, label_list, tokenizer, device)
torch.onnx.export(model,
dummy_input,
onnx_model_file,
input_names=['input_ids', 'input_mask', 'segment_ids'],
output_names=['output_logit'],
verbose=True)
optimized_model = onnx.optimizer.optimize(onnx.load(onnx_model_file),
[pass_ for pass_ in onnx.optimizer.get_available_passes()
if 'split' not in pass_])
optimized_model = onnx.utils.polish_model(optimized_model)
onnx.save(optimized_model, os.path.join(args.output_dir, 'optimized_model.onnx'))
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if args.epochs is None:
args.epochs = 90
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir,
args.verbose)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(
filter(None, [args.compress, args.qe_stats_file
]), # remove both None and empty strings
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
if args.evaluate:
args.deterministic = True
if args.deterministic:
distiller.set_deterministic(
args.seed) # For experiment reproducability
else:
if args.seed is not None:
distiller.set_seed(args.seed)
# Turn on CUDNN benchmark mode for best performance. This is usually "safe" for image
# classification models, as the input sizes don't change during the run
# See here: https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
cudnn.benchmark = True
start_epoch = 0
ending_epoch = args.epochs
perf_scores_history = []
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
raise ValueError(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
raise ValueError(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = distiller.apputils.classification_dataset_str_from_arch(
args.arch)
args.num_classes = distiller.apputils.classification_num_classes(
args.dataset)
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model, config = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# TODO(barrh): args.deprecated_resume is deprecated since v0.3.1
if args.deprecated_resume:
msglogger.warning(
'The "--resume" flag is deprecated. Please use "--resume-from=YOUR_PATH" instead.'
)
if not args.reset_optimizer:
msglogger.warning(
'If you wish to also reset the optimizer, call with: --reset-optimizer'
)
args.reset_optimizer = True
args.resumed_checkpoint_path = args.deprecated_resume
# We can optionally resume from a checkpoint
optimizer = None
if args.resumed_checkpoint_path:
model, compression_scheduler, optimizer, start_epoch = apputils.load_checkpoint(
model, args.resumed_checkpoint_path, model_device=args.device)
elif args.load_model_path:
model = apputils.load_lean_checkpoint(model,
args.load_model_path,
model_device=args.device)
if args.reset_optimizer:
start_epoch = 0
if optimizer is not None:
optimizer = None
msglogger.info(
'\nreset_optimizer flag set: Overriding resumed optimizer and resetting epoch count to 0'
)
# Define loss function (criterion)
if "ssd" in args.arch:
neg_pos_ratio = 3
criterion = MultiboxLoss(config.priors,
iou_threshold=0.5,
neg_pos_ratio=neg_pos_ratio,
center_variance=0.1,
size_variance=0.2,
device=args.device,
reduction="sum",
class_reduction=True,
verbose=0)
else:
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is None:
if "ssd" in args.arch:
base_net_lr = args.lr
extra_layers_lr = args.lr
params = [{
'params': model.base_net.parameters(),
'lr': base_net_lr
}, {
'params':
itertools.chain(model.source_layer_add_ons.parameters(),
model.extras.parameters()),
'lr':
extra_layers_lr
}, {
'params':
itertools.chain(model.regression_headers.parameters(),
model.classification_headers.parameters())
}]
else:
params = model.parameters()
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
for summary in args.summary:
distiller.model_summary(model, summary, args.dataset)
return
if args.export_onnx is not None:
return distiller.export_img_classifier_to_onnx(model,
os.path.join(
msglogger.logdir,
args.export_onnx),
args.dataset,
add_softmax=True,
verbose=False)
if args.qe_calibration:
return acts_quant_stats_collection(model, criterion, pylogger, args)
if args.activation_histograms:
return acts_histogram_collection(model, criterion, pylogger, args)
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = load_data(args, config=config)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if args.resumed_checkpoint_path else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, \
"You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resumed_checkpoint_path.replace(
".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
args.kd_policy = None
if args.kd_teacher:
teacher, _ = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
parallel=not args.load_serialized,
device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
raw_teacher_model_path = msglogger.logdir + "/raw_teacher.pth.tar"
if not os.path.exists(raw_teacher_model_path):
teacher.save(raw_teacher_model_path)
msglogger.info(Fore.CYAN + '\tRaw Teacher Model saved: {0}'.format(
raw_teacher_model_path) + Style.RESET_ALL)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model,
teacher,
args.kd_temp,
dlw,
loss_type=args.kd_loss_type,
focal_alpha=args.kd_focal_alpha,
use_adaptive=args.kd_focal_adaptive,
verbose=0)
compression_scheduler.add_policy(args.kd_policy,
starting_epoch=args.kd_start_epoch,
ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'
.format(start_epoch, ending_epoch))
raise ValueError('Epochs parameter is too low. Nothing to do.')
for epoch in range(start_epoch, ending_epoch):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(
epoch, metrics=(vloss if (epoch != start_epoch) else 10**6))
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint
update_training_scores_history(perf_scores_history, model, top1, top5,
epoch, args.num_best_scores)
is_best = epoch == perf_scores_history[0].epoch
checkpoint_extras = {
'current_top1': top1,
'best_top1': perf_scores_history[0].top1,
'best_epoch': perf_scores_history[0].epoch
}
try:
raw_fullpath_best = apputils.save_checkpoint(
epoch,
args.arch,
model,
optimizer=optimizer,
scheduler=compression_scheduler,
extras=checkpoint_extras,
is_best=is_best,
name=args.name,
dir=msglogger.logdir)
except Exception as ex:
# keep previous fullpath_best
pass
mlflow.log_artifacts(msglogger.logdir)
# Finally run results on the test set
eval_params = {
"model_type": args.arch,
"model_path": raw_fullpath_best,
"dataset_path": args.data,
"label_path": "models/voc-model-labels.txt"
}
mlflow.projects.run(uri=".",
entry_point="eval",
use_conda=False,
parameters=eval_params)
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
apputils.log_execution_env_state(args.compress,
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
best_epochs = list()
if args.deterministic:
if args.loaders is None:
args.loaders = 1
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.loaders > 1:
msglogger.error(
'ERROR: Setting --deterministic requires setting --loaders to 0 or 1'
)
exit(1)
# Use a well-known seed, for repeatability of experiments
distiller.set_deterministic()
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.use_cpu or (args.gpus is None
and not torch.cuda.is_available()) or (args.gpus
== ''):
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
if args.loaders is None:
active_gpus = args.gpus if args.gpus is not None else torch.cuda.device_count(
)
args.loaders = max(parser.DEFAULT_LOADERS_COUNT,
parser.DEFAULT_LOADERS_COUNT * active_gpus)
msglogger.debug('Number of data loaders set to: {}'.format(args.loaders))
# Infer the dataset from the model name
args.dataset = 'cifar10' if 'cifar' in args.arch else 'imagenet'
args.num_classes = 10 if args.dataset == 'cifar10' else 1000
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = create_model(args.pretrained,
args.dataset,
args.arch,
parallel=not args.load_serialized,
device_ids=args.gpus)
compression_scheduler = None
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
optimizer = None
resumed_training_steps = None
if args.resume or args.load_state_dict:
if args.resume and not args.reset_optimizer:
# initiate SGD with dummy lr
optimizer = torch.optim.SGD(model.parameters(), lr=0.36787944117)
model, compression_scheduler, optimizer, start_epoch, resumed_training_steps = apputils.load_checkpoint(
model, args.resume or args.load_state_dict, optimizer=optimizer)
model.to(args.device)
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().to(args.device)
if optimizer is not None:
# optimizer was resumed from checkpoint
# check if user has tried to set optimizer arguments
# if so, ignore arguments with a warning.
optimizer_group_args = [
'lr', 'learning-rate', 'momentum', 'weight-decay', 'wd'
]
user_optim_args = [
x for x in optimizer_group_args for arg in sys.argv
if arg.startswith('--' + x)
]
if user_optim_args:
msglogger.warning(
'{} optimizer arguments are ignored.'.format(user_optim_args))
msglogger.info(
'setting optimizer arguments when optimizer is resumed '
'from checkpoint is forbidden.')
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC:
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy:
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports.
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info(
'\tStats will be collected for {:.1%} of test dataset'.format(
args.qe_calibration))
msglogger.info(
'\tSetting constant seeds and converting model to serialized execution'
)
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(
create_quantization_stats_collector(model))
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
train_loader, val_loader, test_loader, _ = apputils.load_data(
args.dataset, os.path.expanduser(args.data), args.batch_size,
args.loaders, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size,
args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
args.trainset_print_period = parser.getPrintPeriod(
args, len(train_loader.sampler), args.batch_size)
args.validset_print_period = parser.getPrintPeriod(args,
len(val_loader.sampler),
args.batch_size)
args.testset_print_period = parser.getPrintPeriod(args,
len(test_loader.sampler),
args.batch_size)
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
compression_scheduler = distiller.file_config(
model, optimizer, args.compress, compression_scheduler,
(start_epoch - 1) if
(args.resume and not args.reset_optimizer) else None)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resume is not None, "You must use --resume to provide a checkpoint file to thinnify"
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resume.replace(".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
)
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_checkpoint(teacher,
chkpt_file=args.kd_resume)[0]
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(
args.kd_policy, range(args.kd_start_epoch, args.epochs, 1))
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if getattr(compression_scheduler, 'global_policy_end_epoch',
None) is not None:
if compression_scheduler.global_policy_end_epoch >= (start_epoch +
args.epochs):
msglogger.warning(
'scheduler requires at least {} epochs, but only {} are sanctioned'
.format(compression_scheduler.global_policy_end_epoch,
args.epochs))
accumulated_training_steps = resumed_training_steps if resumed_training_steps is not None else 0
for epoch in range(start_epoch, start_epoch + args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
try:
train(train_loader,
model,
criterion,
optimizer,
epoch,
accumulated_training_steps,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
except RuntimeError as e:
if ('cuda out of memory' in str(e).lower()):
msglogger.error(
'CUDA memory failure has been detected.\n'
'Sometimes it helps to decrease batch size.\n'
'e.g. Add the following flag to your call: --batch-size={}'
.format(args.batch_size // 10))
raise
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity:
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
accumulated_training_steps += math.ceil(
len(train_loader.sampler) / train_loader.batch_size)
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Performance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
tflogger.log_training_progress(stats, epoch, None)
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
if getattr(compression_scheduler, 'global_policy_end_epoch',
None) is None or (
compression_scheduler.global_policy_end_epoch <= epoch):
# Update the list of top scores achieved since all policies have concluded
if top1 > 0:
best_epochs.append(
distiller.MutableNamedTuple({
'top1': top1,
'top5': top5,
'epoch': epoch
}))
# Keep best_epochs sorted from best to worst
# Sort by top1 first, secondary sort by top5, and so forth
best_epochs.sort(key=operator.attrgetter('top1', 'top5', 'epoch'),
reverse=True)
for score in best_epochs[:args.num_best_scores]:
msglogger.info('==> Best Top1: %.3f Top5: %.3f on epoch: %d',
score.top1, score.top5, score.epoch)
is_best = best_epochs and (epoch == best_epochs[0].epoch)
apputils.save_checkpoint(epoch, args.arch, model, optimizer,
compression_scheduler,
best_epochs[0].top1 if best_epochs else None,
is_best, args.name, msglogger.logdir,
accumulated_training_steps)
# Finally run results on the test set
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
if args.thinnify:
#zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resumed_checkpoint_path is not None, "You must use --resume-from to provide a checkpoint file to thinnify"
distiller.remove_filters(model, compression_scheduler.zeros_mask_dict, args.arch, args.dataset, optimizer=None)
apputils.save_checkpoint(0, args.arch, model, optimizer=None, scheduler=compression_scheduler,
name="{}_thinned".format(args.resumed_checkpoint_path.replace(".pth.tar", "")), dir=msglogger.logdir)
print("Note: your model may have collapsed to random inference, so you may want to fine-tune")
return
args.kd_policy = None
if args.kd_teacher:
teacher = create_model(args.kd_pretrained, args.dataset, args.kd_teacher, device_ids=args.gpus)
if args.kd_resume:
teacher = apputils.load_lean_checkpoint(teacher, args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt, args.kd_student_wt, args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy, starting_epoch=args.kd_start_epoch, ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
if start_epoch >= ending_epoch:
msglogger.error(
'epoch count is too low, starting epoch is {} but total epochs set to {}'.format(
start_epoch, ending_epoch))
def main():
script_dir = os.path.dirname(__file__)
module_path = os.path.abspath(os.path.join(script_dir, '..', '..'))
global msglogger
# Parse arguments
args = parser.get_parser().parse_args()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
msglogger = apputils.config_pylogger(
os.path.join(script_dir, 'logging.conf'), args.name, args.output_dir)
# Log various details about the execution environment. It is sometimes useful
# to refer to past experiment executions and this information may be useful.
# 记录有关执行环境的各种详细信息。有时是有用的
# 参考过去的实验执行，这些信息可能有用。
apputils.log_execution_env_state(args.compress,
msglogger.logdir,
gitroot=module_path)
msglogger.debug("Distiller: %s", distiller.__version__)
start_epoch = 0
perf_scores_history = []
if args.deterministic:
# Experiment reproducibility is sometimes important. Pete Warden expounded about this
# in his blog: https://petewarden.com/2018/03/19/the-machine-learning-reproducibility-crisis/
# In Pytorch, support for deterministic execution is still a bit clunky.
if args.workers > 1:
msglogger.error(
'ERROR: Setting --deterministic requires setting --workers/-j to 0 or 1'
) # 错误：设置--确定性要求将--workers/-j设置为0或1
exit(1) # 正常退出程序
# Use a well-known seed, for repeatability of experiments 使用一种众所周知的种子，用于实验的重复性。
distiller.set_deterministic()
else:
# This issue: https://github.com/pytorch/pytorch/issues/3659
# Implies that cudnn.benchmark should respect cudnn.deterministic, but empirically we see that
# results are not re-produced when benchmark is set. So enabling only if deterministic mode disabled.
cudnn.benchmark = True
if args.cpu or not torch.cuda.is_available():
# Set GPU index to -1 if using CPU
args.device = 'cpu'
args.gpus = -1
else:
args.device = 'cuda'
if args.gpus is not None:
try:
args.gpus = [int(s) for s in args.gpus.split(',')]
except ValueError:
msglogger.error(
'ERROR: Argument --gpus must be a comma-separated list of integers only'
)
exit(1)
available_gpus = torch.cuda.device_count()
for dev_id in args.gpus:
if dev_id >= available_gpus:
msglogger.error(
'ERROR: GPU device ID {0} requested, but only {1} devices available'
.format(dev_id, available_gpus))
exit(1)
# Set default device in case the first one on the list != 0
torch.cuda.set_device(args.gpus[0])
# Infer the dataset from the model name
args.dataset = 'cousm'
if args.earlyexit_thresholds:
args.num_exits = len(args.earlyexit_thresholds) + 1
args.loss_exits = [0] * args.num_exits
args.losses_exits = []
args.exiterrors = []
# Create the model
model = ResNet152()
# model = torch.nn.DataParallel(model, device_ids=args.gpus) # 并行GPU
model.to(args.device)
compression_scheduler = None # 压缩调度
# Create a couple of logging backends. TensorBoardLogger writes log files in a format
# that can be read by Google's Tensor Board. PythonLogger writes to the Python logger.
# 创建两个日志后端 TensorBoardLogger以Google的Tensor板可以读取的格式写入日志文件。python logger将写入python记录器。
tflogger = TensorBoardLogger(msglogger.logdir)
pylogger = PythonLogger(msglogger)
# capture thresholds for early-exit training
if args.earlyexit_thresholds:
msglogger.info('=> using early-exit threshold values of %s',
args.earlyexit_thresholds)
# We can optionally resume from a checkpoint
if args.resume: # 加载训练模型
# checkpoint = torch.load(args.resume)
# model.load_state_dict(checkpoint['state_dict'])
model, compression_scheduler, start_epoch = apputils.load_checkpoint(
model, chkpt_file=args.resume)
model.to(args.device)
# Define loss function (criterion) and optimizer # 定义损失函数和优化器SGD
criterion = nn.CrossEntropyLoss().to(args.device)
# optimizer = torch.optim.SGD(model.fc.parameters(), lr=args.lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
optimizer = torch.optim.Adam(model.model.fc.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
msglogger.info('Optimizer Type: %s', type(optimizer))
msglogger.info('Optimizer Args: %s', optimizer.defaults)
if args.AMC: # 自动化的深层压缩
return automated_deep_compression(model, criterion, optimizer,
pylogger, args)
if args.greedy: # 贪婪的
return greedy(model, criterion, optimizer, pylogger, args)
# This sample application can be invoked to produce various summary reports. # 可以调用此示例应用程序来生成各种摘要报告。
if args.summary:
return summarize_model(model, args.dataset, which_summary=args.summary)
# 激活统计收集器
activations_collectors = create_activation_stats_collectors(
model, *args.activation_stats)
if args.qe_calibration:
msglogger.info('Quantization calibration stats collection enabled:')
msglogger.info(
'\tStats will be collected for {:.1%} of test dataset'.format(
args.qe_calibration))
msglogger.info(
'\tSetting constant seeds and converting model to serialized execution'
)
distiller.set_deterministic()
model = distiller.make_non_parallel_copy(model)
activations_collectors.update(
create_quantization_stats_collector(model)) # 量化统计收集器
args.evaluate = True
args.effective_test_size = args.qe_calibration
# Load the datasets: the dataset to load is inferred from the model name passed
# in args.arch. The default dataset is ImageNet, but if args.arch contains the
# substring "_cifar", then cifar10 is used.
# 加载数据集：从传递的模型名称推断要加载的数据集
train_loader, val_loader, test_loader, _ = get_data_loaders(
datasets_fn, r'/home/tian/Desktop/image_yasuo', args.batch_size,
args.workers, args.validation_split, args.deterministic,
args.effective_train_size, args.effective_valid_size,
args.effective_test_size)
msglogger.info('Dataset sizes:\n\ttraining=%d\n\tvalidation=%d\n\ttest=%d',
len(train_loader.sampler), len(val_loader.sampler),
len(test_loader.sampler))
# 可以调用此示例应用程序来对模型执行敏感性分析。输出保存到csv和png。
if args.sensitivity is not None:
sensitivities = np.arange(args.sensitivity_range[0],
args.sensitivity_range[1],
args.sensitivity_range[2])
return sensitivity_analysis(model, criterion, test_loader, pylogger,
args, sensitivities)
if args.evaluate:
return evaluate_model(model, criterion, test_loader, pylogger,
activations_collectors, args,
compression_scheduler)
if args.compress:
# The main use-case for this sample application is CNN compression. Compression
# requires a compression schedule configuration file in YAML.
# #这个示例应用程序的主要用例是CNN压缩
# #需要yaml中的压缩计划配置文件。
compression_scheduler = distiller.file_config(model, optimizer,
args.compress,
compression_scheduler)
# Model is re-transferred to GPU in case parameters were added (e.g. PACTQuantizer)
# 如果添加了参数（如PactQualifier），则模型会重新传输到GPU。
model.to(args.device)
elif compression_scheduler is None:
compression_scheduler = distiller.CompressionScheduler(model) # 压缩计划程序
if args.thinnify:
# zeros_mask_dict = distiller.create_model_masks_dict(model)
assert args.resume is not None, "You must use --resume to provide a checkpoint file to thinnify" # 必须使用--resume提供检查点文件以细化
distiller.remove_filters(model,
compression_scheduler.zeros_mask_dict,
args.arch,
args.dataset,
optimizer=None)
apputils.save_checkpoint(0,
args.arch,
model,
optimizer=None,
scheduler=compression_scheduler,
name="{}_thinned".format(
args.resume.replace(".pth.tar", "")),
dir=msglogger.logdir)
print(
"Note: your model may have collapsed to random inference, so you may want to fine-tune"
) # 注意：您的模型可能已折叠为随机推理，因此您可能需要对其进行微调。
return
args.kd_policy = None # 蒸馏
if args.kd_teacher:
teacher = create_model(args.kd_pretrained,
args.dataset,
args.kd_teacher,
device_ids=args.gpus)
if args.kd_resume:
teacher, _, _ = apputils.load_checkpoint(teacher,
chkpt_file=args.kd_resume)
dlw = distiller.DistillationLossWeights(args.kd_distill_wt,
args.kd_student_wt,
args.kd_teacher_wt)
args.kd_policy = distiller.KnowledgeDistillationPolicy(
model, teacher, args.kd_temp, dlw)
compression_scheduler.add_policy(args.kd_policy,
starting_epoch=args.kd_start_epoch,
ending_epoch=args.epochs,
frequency=1)
msglogger.info('\nStudent-Teacher knowledge distillation enabled:')
msglogger.info('\tTeacher Model: %s', args.kd_teacher)
msglogger.info('\tTemperature: %s', args.kd_temp)
msglogger.info('\tLoss Weights (distillation | student | teacher): %s',
' | '.join(['{:.2f}'.format(val) for val in dlw]))
msglogger.info('\tStarting from Epoch: %s', args.kd_start_epoch)
lr = args.lr
lr_decay = 0.5
for epoch in range(start_epoch, args.epochs):
# This is the main training loop.
msglogger.info('\n')
if compression_scheduler:
compression_scheduler.on_epoch_begin(epoch)
# Train for one epoch
with collectors_context(activations_collectors["train"]) as collectors:
train(train_loader,
model,
criterion,
optimizer,
epoch,
compression_scheduler,
loggers=[tflogger, pylogger],
args=args)
distiller.log_weights_sparsity(model,
epoch,
loggers=[tflogger, pylogger])
distiller.log_activation_statsitics(
epoch,
"train",
loggers=[tflogger],
collector=collectors["sparsity"])
if args.masks_sparsity: # 打印掩盖稀疏表 在end of each epoch
msglogger.info(
distiller.masks_sparsity_tbl_summary(
model, compression_scheduler))
# evaluate on validation set
with collectors_context(activations_collectors["valid"]) as collectors:
top1, top5, vloss = validate(val_loader, model, criterion,
[pylogger], args, epoch)
distiller.log_activation_statsitics(
epoch,
"valid",
loggers=[tflogger],
collector=collectors["sparsity"])
save_collectors_data(collectors, msglogger.logdir)
stats = ('Peformance/Validation/',
OrderedDict([('Loss', vloss), ('Top1', top1),
('Top5', top5)]))
distiller.log_training_progress(stats,
None,
epoch,
steps_completed=0,
total_steps=1,
log_freq=1,
loggers=[tflogger])
if compression_scheduler:
compression_scheduler.on_epoch_end(epoch, optimizer)
# Update the list of top scores achieved so far, and save the checkpoint # 更新到目前为止获得的最高分数列表，并保存检查点
sparsity = distiller.model_sparsity(model)
perf_scores_history.append(
distiller.MutableNamedTuple({
'sparsity': sparsity,
'top1': top1,
'top5': top5,
'epoch': epoch
}))
# Keep perf_scores_history sorted from best to worst
# Sort by sparsity as main sort key, then sort by top1, top5 and epoch
# 保持绩效分数历史记录从最好到最差的排序
# 按稀疏度排序为主排序键，然后按top1、top5、epoch排序
perf_scores_history.sort(key=operator.attrgetter(
'sparsity', 'top1', 'top5', 'epoch'),
reverse=True)
for score in perf_scores_history[:args.num_best_scores]:
msglogger.info(
'==> Best [Top1: %.3f Top5: %.3f Sparsity: %.2f on epoch: %d]',
score.top1, score.top5, score.sparsity, score.epoch)
is_best = epoch == perf_scores_history[0].epoch
apputils.save_checkpoint(epoch, args.arch, model, optimizer,
compression_scheduler,
perf_scores_history[0].top1, is_best,
args.name, msglogger.logdir)
if not is_best:
lr = lr * lr_decay
# 当loss大于上一次loss,降低学习率
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# Finally run results on the test set # 最后在测试集上运行结果
test(test_loader,
model,
criterion, [pylogger],
activations_collectors,
args=args)
