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))
