def train_auto_compressor(self):
using_fm_reconstruction = self.args.amc_prune_method == 'fm-reconstruction'
fixed_subset, sequential = (using_fm_reconstruction,
using_fm_reconstruction)
msglogger.info("AMC: fixed_subset=%s\tsequential=%s" %
(fixed_subset, sequential))
train_loader, val_loader, test_loader = classifier.load_data(
self.args, fixed_subset, sequential)
self.args.display_confusion = False
validate_fn = partial(classifier.test,
test_loader=val_loader,
criterion=self.criterion,
loggers=self.pylogger,
args=self.args,
activations_collectors=None)
train_fn = partial(classifier.train,
train_loader=train_loader,
criterion=self.criterion,
loggers=self.pylogger,
args=self.args)
save_checkpoint_fn = partial(apputils.save_checkpoint,
arch=self.args.arch,
dir=msglogger.logdir)
optimizer_data = {
'lr': self.args.lr,
'momentum': self.args.momentum,
'weight_decay': self.args.weight_decay
}
return train_auto_compressor(self.model, self.args, optimizer_data,
validate_fn, save_checkpoint_fn, train_fn)
def greedy(model, criterion, optimizer, loggers, args):
train_loader, val_loader, test_loader = classifier.load_data(args)
test_fn = partial(classifier.test, test_loader=test_loader, criterion=criterion,
loggers=loggers, args=args, activations_collectors=None)
train_fn = partial(classifier.train, train_loader=train_loader, criterion=criterion, args=args)
assert args.greedy_target_density is not None
distiller.pruning.greedy_filter_pruning.greedy_pruner(model, args,
args.greedy_target_density,
args.greedy_pruning_step,
test_fn, train_fn)
def load_test_data(args):
test_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True)
return test_loader
q_weight_name = module_name + ".weight"
if q_weight_name == name_in:
return module
if __name__ == "__main__":
print("START")
args = get_default_args()
# args.epochs = float('inf') # hack for args parsing so there's no error in epochs
cc = classifier.ClassifierCompressor(args,
script_dir=os.path.dirname(__file__))
args = deepcopy(
cc.args
) # Get back args after modifications in ClassifierCompressor.__init__
eval_data_loader = classifier.load_data(args,
load_train=False,
load_val=False,
load_test=True)
# logging
logging.getLogger().setLevel(logging.WARNING)
msglogger = logging.getLogger(__name__)
msglogger.setLevel(logging.INFO)
def test_fn(model):
top1, top5, losses = classifier.test(eval_data_loader, model,
cc.criterion,
[cc.tflogger, cc.pylogger], None,
args)
# pdb.set_trace()
return top1, top5, losses
def finetune_directory(ft_dir,
stats_file,
app_args,
cleanup_ft_dir=False,
checkpoints=None):
"""Fine tune all the checkpoint files we find in the immediate-directory specified.
For each checkpoint file we find, we create and queue a FinetuningTask.
A FinetuningProcess will pickup the FinetuningTask and process it.
"""
print("Fine-tuning directory %s" % ft_dir)
if not checkpoints:
# Get a list of the checkpoint files
checkpoints = glob.glob(os.path.join(ft_dir, "*checkpoint.pth.tar"))
assert checkpoints
# We create a subdirectory, where we will write all of our output
ft_output_dir = os.path.join(ft_dir, "ft")
os.makedirs(ft_output_dir, exist_ok=True)
print("Writing results to directory %s" % ft_output_dir)
app_args.output_dir = ft_output_dir
# Multi-process queues
tasks = multiprocessing.JoinableQueue()
results = multiprocessing.Queue()
# Create and launch the fine-tuning processes
processes = []
n_processes = min(app_args.processes, len(checkpoints))
for i in range(n_processes):
# Pre-load the data-loaders of each fine-tuning process once
app = classifier.ClassifierCompressor(
app_args, script_dir=os.path.dirname(__file__))
data_loader = classifier.load_data(app.args)
# Delete log directories
shutil.rmtree(app.logdir)
processes.append(FinetuningProcess(tasks, results, data_loader))
# Start the process
processes[-1].start()
n_gpus = torch.cuda.device_count()
# Enqueue all of the fine-tuning tasks
for (instance, ckpt_file) in enumerate(checkpoints):
tasks.put(FinetuningTask(args=(ckpt_file, instance % n_gpus,
app_args)))
# Push an end-of-tasks marker
for i in range(len(processes)):
tasks.put(None)
# Wait until all tasks finish
tasks.join()
# Start printing results
results_dict = OrderedDict()
while not results.empty():
result = results.get()
results_dict[result[0]] = result[1]
# Read the results of the AMC experiment (we'll want to use some of the data)
import pandas as pd
df = pd.read_csv(os.path.join(ft_dir, "amc.csv"))
assert len(results_dict) > 0
# Log some info for each checkpoint
for ckpt_name in sorted(results_dict.keys()):
net_search_results = df[df["ckpt_name"] ==
ckpt_name[:-len("_checkpoint.pth.tar")]]
search_top1 = net_search_results["top1"].iloc[0]
normalized_macs = net_search_results["normalized_macs"].iloc[0]
log_entry = (ft_output_dir, ckpt_name, normalized_macs, search_top1,
*results_dict[ckpt_name])
print("%s <> %s: %.2f %.2f %.2f %.2f %.2f" % log_entry)
stats_file.add_record(log_entry)
if cleanup_ft_dir:
# cleanup: remove the "ft" directory
shutil.rmtree(ft_output_dir)
msglogger.info('best_overrides_dict: %s' % best_overrides_dict)
msglogger.info('Best score: %f' % eval_fn(quantizer.model))
return model, best_overrides_dict
if __name__ == "__main__":
args = get_default_args()
args.epochs = float(
'inf') # hack for args parsing so there's no error in epochs
cc = classifier.ClassifierCompressor(args,
script_dir=os.path.dirname(__file__))
args = deepcopy(
cc.args
) # Get back args after modifications in ClassifierCompressor.__init__
eval_data_loader = classifier.load_data(args,
load_train=False,
load_val=False)
# quant calibration dataloader:
args.effective_test_size = args.qe_calib_portion
args.batch_size = args.qe_calib_batchsize
calib_data_loader = classifier.load_data(args,
fixed_subset=True,
load_train=False,
load_val=False)
# logging
logging.getLogger().setLevel(logging.WARNING)
msglogger = logging.getLogger(__name__)
msglogger.setLevel(logging.INFO)
def test_fn(model):
def image_classifier_ptq_lapq(model, criterion, loggers, args):
# data loader function for splitting the validation set.
args = deepcopy(args)
effective_test_size_bak = args.effective_test_size
args.effective_test_size = args.lapq_eval_size
eval_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True, fixed_subset=True)
args.effective_test_size = effective_test_size_bak
test_data_loader = classifier.load_data(args, load_train=False, load_val=False, load_test=True)
model = model.eval()
device = next(model.parameters()).device
if args.lapq_eval_memoize_dataloader:
images_batches = []
targets_batches = []
for images, targets in eval_data_loader:
images_batches.append(images.to(device))
targets_batches.append(targets.to(device))
memoized_data_loader = [(torch.cat(images_batches), torch.cat(targets_batches))]
else:
memoized_data_loader = None
def eval_fn(model):
if memoized_data_loader:
loss = 0
for images, targets in memoized_data_loader:
outputs = model(images)
loss += criterion(outputs, targets).item()
loss = loss / len(memoized_data_loader)
else:
_, _, loss = classifier.test(eval_data_loader, model, criterion, loggers, None, args)
return loss
def test_fn(model):
top1, top5, loss = classifier.test(test_data_loader, model, criterion, loggers, None, args)
return OrderedDict([('top-1', top1), ('top-5', top5), ('loss', loss)])
args.device = device
if args.resumed_checkpoint_path:
args.load_model_path = args.resumed_checkpoint_path
if args.load_model_path:
msglogger.info("Loading checkpoint from %s" % args.load_model_path)
model = apputils.load_lean_checkpoint(model, args.load_model_path,
model_device=args.device)
quantizer = distiller.quantization.PostTrainLinearQuantizer.from_args(model, args)
dummy_input = torch.rand(*model.input_shape, device=args.device)
model, qp_dict = lapq.ptq_coordinate_search(quantizer, dummy_input, eval_fn, test_fn=test_fn,
**lapq.cmdline_args_to_dict(args))
results = test_fn(quantizer.model)
msglogger.info("Arch: %s \tTest: \t top1 = %.3f \t top5 = %.3f \t loss = %.3f" %
(args.arch, results['top-1'], results['top-5'], results['loss']))
distiller.yaml_ordered_save('%s.quant_params_dict.yaml' % args.arch, qp_dict)
distiller.apputils.save_checkpoint(0, args.arch, model,
extras={'top1': results['top-1'], 'qp_dict': qp_dict}, name=args.name,
dir=msglogger.logdir)
if __name__ == "__main__":
args = get_default_args()
args.epochs = float(
'inf') # hack for args parsing so there's no error in epochs
cc = classifier.ClassifierCompressor(args,
script_dir=os.path.dirname(__file__))
args = deepcopy(cc.args)
effective_test_size_bak = args.effective_test_size
args.effective_test_size = args.opt_val_size
eval_data_loader = classifier.load_data(args,
load_train=False,
load_val=False,
load_test=True,
fixed_subset=True)
args.effective_test_size = effective_test_size_bak
test_data_loader = classifier.load_data(args,
load_train=False,
load_val=False,
load_test=True)
# logging
logging.getLogger().setLevel(logging.WARNING)
msglogger = logging.getLogger(__name__)
msglogger.setLevel(logging.INFO)
model = cc.model.eval()