def _log_model_info(self): msglogger.debug("Model %s has %d modules (%d pruned)", self.app_args.arch, self.net_wrapper.model_metadata.num_layers(), self.net_wrapper.model_metadata.num_pruned_layers()) msglogger.debug("\tTotal MACs: %s" % distiller.pretty_int(self.original_model_macs)) msglogger.debug("\tTotal weights: %s" % distiller.pretty_int(self.original_model_size))
def __init__(self, model, dataset, arch, data_loader, validate_fn, save_checkpoint_fn): self.pylogger = distiller.data_loggers.PythonLogger(msglogger) self.tflogger = distiller.data_loggers.TensorBoardLogger( msglogger.logdir) self.action_space = RandomADCActionSpace() self.dataset = dataset self.arch = arch self.data_loader = data_loader self.validate_fn = validate_fn self.save_checkpoint_fn = save_checkpoint_fn self.orig_model = model self.conv_layers, self.dense_model_macs = collect_conv_details( model, dataset) self.reset(init_only=True) msglogger.info("Model %s has %d Convolution layers", arch, len(self.conv_layers)) msglogger.info("\tTotal MACs: %s" % distiller.pretty_int(self.dense_model_macs)) self.debug_stats = {'episode': 0} # Gym # spaces documentation: https://gym.openai.com/docs/ self.action_space = spaces.Box(0, 1, shape=(1, )) self.observation_space = spaces.Box(0, float("inf"), shape=(self.STATE_EMBEDDING_LEN, ))
def __init__(self, model, app_args, amc_cfg, services): self.pylogger = distiller.data_loggers.PythonLogger(msglogger) self.tflogger = distiller.data_loggers.TensorBoardLogger(msglogger.logdir) self.orig_model = model self.app_args = app_args self.amc_cfg = amc_cfg self.services = services self.conv_layers, self.dense_model_macs, self.dense_model_size = collect_conv_details(model, self.app_args.dataset) self.reset(init_only=True) msglogger.info("Model %s has %d Convolution layers", self.app_args.arch, len(self.conv_layers)) msglogger.info("\tTotal MACs: %s" % distiller.pretty_int(self.dense_model_macs)) msglogger.info("Configuration:\n\tonehot_encoding={}\n\tnormalize_obs={}".format(self.amc_cfg.onehot_encoding, self.amc_cfg.normalize_obs)) self.debug_stats = {'episode': 0} self.action_low = amc_cfg.action_range[0] self.action_high = amc_cfg.action_range[1] # Gym spaces documentation: https://gym.openai.com/docs/ self.action_space = spaces.Box(self.action_low, self.action_high, shape=(1,)) self.action_space.default_action = self.action_low self.STATE_EMBEDDING_LEN = len(Observation._fields) if self.amc_cfg.onehot_encoding: self.STATE_EMBEDDING_LEN += (self.amc_cfg.conv_cnt - 1) self.observation_space = spaces.Box(0, float("inf"), shape=(self.STATE_EMBEDDING_LEN,)) fields = ['top1', 'reward', 'total_macs', 'normalize_macs', 'total_nnz'] with open(r'amc.csv', 'w') as f: writer = csv.writer(f) writer.writerow(fields)
def log_training_progress(self, stats_dict, epoch, completed, total, freq): stats_dict = stats_dict[1] if epoch > -1: log = 'Epoch: [{}][{:5d}/{:5d}] '.format(epoch, completed, int(total)) else: log = 'Test: [{:5d}/{:5d}] '.format(completed, int(total)) for name, val in stats_dict.items(): if isinstance(val, int): log += '{name} {val} '.format(name=name, val=distiller.pretty_int(val)) else: log += '{name} {val:.6f} '.format(name=name, val=val) self.pylogger.info(log)
def __init__(self, model, dataset, arch, data_loader, validate_fn, save_checkpoint_fn, action_range, onehot_encoding, normalize_obs, desired_reduction, reward_fn): self.pylogger = distiller.data_loggers.PythonLogger(msglogger) self.tflogger = distiller.data_loggers.TensorBoardLogger( msglogger.logdir) self.dataset = dataset self.arch = arch self.data_loader = data_loader self.validate_fn = validate_fn self.save_checkpoint_fn = save_checkpoint_fn self.orig_model = model self.onehot_encoding = onehot_encoding self.normalize_obs = normalize_obs self.max_reward = -1000 self.reward_fn = reward_fn self.conv_layers, self.dense_model_macs, self.dense_model_size = collect_conv_details( model, dataset) self.reset(init_only=True) msglogger.info("Model %s has %d Convolution layers", arch, len(self.conv_layers)) msglogger.info("\tTotal MACs: %s" % distiller.pretty_int(self.dense_model_macs)) msglogger.info( "Configuration:\n\tonehot_encoding={}\n\tnormalize_obs={}".format( self.onehot_encoding, self.normalize_obs)) self.debug_stats = {'episode': 0} self.action_low = action_range[0] self.action_high = action_range[1] # Gym # spaces documentation: https://gym.openai.com/docs/ self.action_space = spaces.Box(self.action_low, self.action_high, shape=(1, )) self.action_space.default_action = self.action_low self.desired_reduction = desired_reduction self.STATE_EMBEDDING_LEN = len(Observation._fields) if self.onehot_encoding: self.STATE_EMBEDDING_LEN += (steps_per_episode - 1) self.observation_space = spaces.Box(0, float("inf"), shape=(self.STATE_EMBEDDING_LEN, )) fields = ['top1', 'reward', 'total_macs', 'total_nnz'] with open(r'amc.csv', 'w') as f: writer = csv.writer(f) writer.writerow(fields)
def log_training_progress(self, stats_dict, epoch, absolute_train_step, train_step_in_epoch, train_steps_per_epoch, **kwargs): stats_dict = stats_dict[1] if epoch > -1: log = 'Epoch: [{}][{:5d}/{:5d}] '.format( epoch, train_step_in_epoch, train_steps_per_epoch) else: log = 'Test: [{:5d}/{:5d}] '.format(train_step_in_epoch, train_steps_per_epoch) for name, val in stats_dict.items(): if isinstance(val, int): log = log + '{name} {val} '.format( name=name, val=distiller.pretty_int(val)) else: log = log + '{name} {val:.6f} '.format(name=name, val=val) self.pylogger.info(log)
def __init__(self, model, app_args, amc_cfg, services): self.pylogger = distiller.data_loggers.PythonLogger(msglogger) self.tflogger = distiller.data_loggers.TensorBoardLogger( msglogger.logdir) self.orig_model = model self.app_args = app_args self.amc_cfg = amc_cfg self.services = services self.net_wrapper = NetworkWrapper(model, app_args, services) self.dense_model_macs, self.dense_model_size = self.net_wrapper.get_model_resources_requirements( model) self.reset(init_only=True) msglogger.info("Model %s has %d Convolution layers", self.app_args.arch, self.net_wrapper.num_layers()) msglogger.info("\tTotal MACs: %s" % distiller.pretty_int(self.dense_model_macs)) log_amc_config(amc_cfg) self.episode = 0 self.best_reward = -1000 self.action_low = amc_cfg.action_range[0] self.action_high = amc_cfg.action_range[1] # Gym spaces documentation: https://gym.openai.com/docs/ if is_using_continuous_action_space(self.amc_cfg.agent_algo): self.action_space = spaces.Box(self.action_low, self.action_high, shape=(1, )) self.action_space.default_action = self.action_low else: self.action_space = spaces.Discrete(10) self.STATE_EMBEDDING_LEN = len(Observation._fields) #self.observation_space = spaces.Box(0, float("inf"), shape=(self.STATE_EMBEDDING_LEN+self.num_layers(),)) self.observation_space = spaces.Box(0, float("inf"), shape=(self.STATE_EMBEDDING_LEN + 1, )) #self.observation_space = spaces.Box(0, float("inf"), shape=(LayerDescLen * self.num_layers(), )) #self.create_network_record_file() self.stats_file = AMCStatsFile( os.path.join(msglogger.logdir, 'amc.csv')) self.ft_stats_file = FineTuneStatsFile( os.path.join(msglogger.logdir, 'ft_top1.csv'))
def __init__(self, model, app_args, amc_cfg, services): self.pylogger = distiller.data_loggers.PythonLogger(msglogger) logdir = logging.getLogger().logdir self.tflogger = distiller.data_loggers.TensorBoardLogger(logdir) self.verbose = False self.orig_model = copy.deepcopy(model) self.app_args = app_args self.amc_cfg = amc_cfg self.services = services try: modules_list = amc_cfg.modules_dict[app_args.arch] except KeyError: msglogger.warning( "!!! The config file does not specify the modules to compress for %s" % app_args.arch) # Default to using all convolution layers distiller.assign_layer_fq_names(model) modules_list = [ mod.distiller_name for mod in model.modules() if type(mod) == torch.nn.Conv2d ] msglogger.warning("Using the following layers: %s" % ", ".join(modules_list)) self.net_wrapper = NetworkWrapper(model, app_args, services, modules_list, amc_cfg.pruning_pattern) self.original_model_macs, self.original_model_size = self.net_wrapper.get_resources_requirements( ) self.reset(init_only=True) msglogger.debug("Model %s has %d modules (%d pruned)", self.app_args.arch, self.net_wrapper.model_metadata.num_layers(), self.net_wrapper.model_metadata.num_pruned_layers()) msglogger.debug("\tTotal MACs: %s" % distiller.pretty_int(self.original_model_macs)) msglogger.debug("\tTotal weights: %s" % distiller.pretty_int(self.original_model_size)) self._max_episode_steps = self.net_wrapper.model_metadata.num_pruned_layers( ) # Hack for Coach-TD3 log_amc_config(amc_cfg) self.episode = 0 self.best_reward = float("-inf") self.action_low = amc_cfg.action_range[0] self.action_high = amc_cfg.action_range[1] if is_using_continuous_action_space(self.amc_cfg.agent_algo): if self.amc_cfg.agent_algo == "ClippedPPO-continuous": self.action_space = spaces.Box(PPO_MIN, PPO_MAX, shape=(1, )) else: self.action_space = spaces.Box(self.action_low, self.action_high, shape=(1, )) self.action_space.default_action = self.action_low else: self.action_space = spaces.Discrete(10) self.observation_space = spaces.Box(0, float("inf"), shape=(len(Observation._fields), )) self.stats_logger = AMCStatsLogger(os.path.join(logdir, 'amc.csv')) self.ft_stats_logger = FineTuneStatsLogger( os.path.join(logdir, 'ft_top1.csv')) if self.amc_cfg.pruning_method == "fm-reconstruction": if self.amc_cfg.pruning_pattern != "channels": raise ValueError( "Feature-map reconstruction is only supported when pruning weights channels" ) from functools import partial def acceptance_criterion(m, mod_names): # Collect feature-maps only for Conv2d layers, if they are in our modules list. return isinstance( m, torch.nn.Conv2d) and m.distiller_name in mod_names # For feature-map reconstruction we need to collect a representative set # of inter-layer feature-maps from distiller.pruning import FMReconstructionChannelPruner collect_intermediate_featuremap_samples( self.net_wrapper.model, self.net_wrapper.validate, partial(acceptance_criterion, mod_names=modules_list), partial( FMReconstructionChannelPruner.cache_featuremaps_fwd_hook, n_points_per_fm=self.amc_cfg.n_points_per_fm))