from __future__ import absolute_import, division, print_function import mxnet as mx import mxnet.ndarray as nd import numpy import os import pickle from collections import OrderedDict import logging from utils import * logger = logging.getLogger(__name__) class Base(object): """Basic wrapper for the symbols Parameters ---------- data_shapes : dict The shapes of tensor variables sym_gen : mx.sym.Symbol Symbol of the network params : None or dict, optional params_grad : None or dict, optional aux_states: initializer: ctx: name: """ def __init__(self, data_shapes, sym_gen, params=None, aux_states=None, default_bucket_kwargs=None, learn_init_keys=None, initializer=mx.init.Xavier(factor_type="in", rnd_type="gaussian", magnitude=2), ctx=mx.gpu(), name='Net'): self.sym_gen = sym_gen bucket_kwargs = default_bucket_kwargs.copy() if \ default_bucket_kwargs is not None else dict() self.curr_bucket_key = None self.ctx = ctx self.name = name self.initializer = initializer if params is None: self.params = None self.params_grad = None else: self.params = OrderedDict([(k, v.copyto(ctx)) for k, v in params.items()]) self.params_grad = OrderedDict([(n, nd.empty(v.shape, ctx=ctx)) for n, v in self.params.items()]) if aux_states is not None: self.aux_states = OrderedDict([(k, v.copyto(ctx)) for k, v in aux_states.items()]) else: self.aux_states = None self._buckets = dict() self.learn_init_keys = learn_init_keys if learn_init_keys is not None else [] self.learn_init_key_shapes = {k: data_shapes[k] for k in self.learn_init_keys} self.switch_bucket(bucket_kwargs=bucket_kwargs, data_shapes=data_shapes) self.acc_grad = None @property def exe(self): """Get the current executor Returns ------- exe : mxnet.executor.Executor """ return self._buckets[self.curr_bucket_key]['exe'][tuple(self.data_shapes.items())] @property def data_shapes(self): return self._buckets[self.curr_bucket_key]['data_shapes'] @property def sym(self): return self._buckets[self.curr_bucket_key]['sym'] def switch_bucket(self, bucket_kwargs=None, data_shapes=None): if bucket_kwargs is not None: self.curr_bucket_key = get_bucket_key(bucket_kwargs=bucket_kwargs) # 1. Check if bucket key exists if self.curr_bucket_key in self._buckets: if data_shapes is not None: if tuple(data_shapes.items()) not in self._buckets[self.curr_bucket_key]['exe']: #TODO Optimize the reshaping functionality! self._buckets[self.curr_bucket_key]['exe'][tuple(data_shapes.items())] = \ self.exe.reshape(partial_shaping=True, allow_up_sizing=True, **data_shapes) self._buckets[self.curr_bucket_key]['data_shapes'] = data_shapes else: self._buckets[self.curr_bucket_key]['data_shapes'] = data_shapes return # 2. If the bucket key does not exist, create new symbol + executor assert data_shapes is not None, "Must set data_shapes for new bucket!" if isinstance(self.sym_gen, mx.symbol.Symbol): sym = self.sym_gen else: sym = self.sym_gen(**dict(self.curr_bucket_key)) arg_names = sym.list_arguments() aux_names = sym.list_auxiliary_states() param_names = [n for n in arg_names if n in self.learn_init_keys or (n not in data_shapes.keys())] for k, v in data_shapes.items(): assert isinstance(v, tuple), "Data_shapes must be tuple! Find k=%s, v=%s, " \ "data_shapes=%s" % (k, str(v), str(data_shapes)) arg_shapes, _, aux_shapes = sym.infer_shape(**data_shapes) arg_name_shape = OrderedDict([(k, s) for k, s in zip(arg_names, arg_shapes)]) if self.params is None: self.params = OrderedDict([(n, nd.empty(arg_name_shape[n], ctx=self.ctx)) for n in param_names]) self.params_grad = OrderedDict([(n, nd.empty(arg_name_shape[n], ctx=self.ctx)) for n in param_names]) if len(self.params) > 0: assert self.initializer is not None, \ 'We must set the initializer if we donnot initialize' \ 'manually the free parameters of the network!!' for k, v in self.params.items(): self.initializer(k, v) else: assert set(arg_name_shape.items()) == \ set(data_shapes.items() + [(k, v.shape) for k, v in self.params.items()]) if self.aux_states is None: self.aux_states = OrderedDict([(k, nd.empty(s, ctx=self.ctx)) for k, s in zip(aux_names, aux_shapes)]) data_inputs = {k: mx.nd.empty(data_shapes[k], ctx=self.ctx) for k in set(data_shapes.keys()) - set(self.learn_init_keys)} if len(self._buckets) > 0: shared_exe = list(list(self._buckets.values())[0]['exe'].values())[0] else: shared_exe = None self._buckets[self.curr_bucket_key] = { 'exe': {tuple(data_shapes.items()): sym.bind(ctx=self.ctx, args=dict(self.params, **data_inputs), args_grad=dict(self.params_grad.items()), aux_states=self.aux_states, shared_exec=shared_exe) }, 'data_shapes': data_shapes, 'sym': sym } def save_params(self, dir_path="", epoch=None): param_saving_path = save_params(dir_path=dir_path, name=self.name, epoch=epoch, params=self.params, aux_states=self.aux_states) misc_saving_path = save_misc(dir_path=dir_path, epoch=epoch, name=self.name, content={'data_shapes': {k: map(int, v) for k, v in self.data_shapes.items()}}) logging.info('Saving %s, params: \"%s\", misc: \"%s\"', self.name, param_saving_path, misc_saving_path) def load_params(self, name="", dir_path="", epoch=None): params, aux_states, param_loading_path = load_params(dir_path=dir_path, epoch=epoch, name=name) logging.info('Loading params from \"%s\" to %s' % (param_loading_path, self.name)) for k, v in params.items(): if k in self.params: logging.debug(' Loading %s %s' %(k, str(v.shape))) self.params[k][:] = v else: logging.warn("Found unused param in the saved model file: %s" % k) for k, v in aux_states.items(): self.aux_states[k][:] = v @property def internal_sym_names(self): return self.sym.get_internals().list_outputs() @property def output_keys(self): return self.sym.list_outputs() def compute_internal(self, sym_name, bucket_kwargs=None, **arg_dict): """ View the internal symbols using the forward function. :param sym_name: :param bucket_kwargs: :param input_dict: :return: """ data_shapes = {k: v.shape for k, v in arg_dict.items()} self.switch_bucket(bucket_kwargs=bucket_kwargs, data_shapes=data_shapes) internal_sym = self.sym.get_internals()[sym_name] data_inputs = {k: mx.nd.empty(v, ctx=self.ctx) for k, v in self.data_shapes.items() if k in internal_sym.list_arguments()} params = {k: v for k, v in self.params.items() if k in internal_sym.list_arguments()} aux_states = {k: v for k, v in self.aux_states.items() if k in internal_sym.list_auxiliary_states()} exe = internal_sym.bind(ctx=self.ctx, args=dict(params, **data_inputs), args_grad=None, grad_req='null', aux_states=aux_states, shared_exec=self.exe) for k, v in arg_dict.items(): exe.arg_dict[k][:] = v exe.forward(is_train=False) assert 1 == len(exe.outputs) for output in exe.outputs: output.wait_to_read() return exe.outputs[0] def forward(self, is_train=False, bucket_kwargs=None, **arg_dict): #import time #start = time.time() data_shapes = {k: v.shape for k, v in arg_dict.items()} for name in self.learn_init_keys: data_shapes[name] = self.learn_init_key_shapes[name] self.switch_bucket(bucket_kwargs=bucket_kwargs, data_shapes=data_shapes) #end = time.time() #print 'Swith Bucket:', end - start #start = time.time() for k, v in arg_dict.items(): assert self.exe.arg_dict[k].shape == v.shape,\ "Shape not match: key %s, need %s, received %s" \ %(k, str(self.exe.arg_dict[k].shape), str(v.shape)) self.exe.arg_dict[k][:] = v self.exe.forward(is_train=is_train) for output in self.exe.outputs: output.wait_to_read() #end = time.time() #print 'Forward:', end - start return self.exe.outputs def backward(self, out_grads=None, **arg_dict): for k, v in arg_dict.items(): assert self.exe.arg_dict[k].shape == v.shape, \ "Shape not match: key %s, need %s, received %s" \ % (k, str(self.exe.arg_dict[k].shape), str(v.shape)) self.exe.arg_dict[k][:] = v self.exe.backward(out_grads=out_grads) def forward_backward(self, bucket_kwargs=None, out_grads=None, **arg_dict): data_shapes = {k: v.shape for k, v in arg_dict.items()} for name in self.learn_init_keys: data_shapes[name] = self.learn_init_key_shapes[name] self.switch_bucket(bucket_kwargs=bucket_kwargs, data_shapes=data_shapes) for k, v in arg_dict.items(): self.exe.arg_dict[k][:] = v self.exe.forward(is_train=True) self.exe.backward(out_grads=out_grads) for output in self.exe.outputs: output.wait_to_read() return self.exe.outputs def update(self, updater, params_grad=None): if params_grad is None: params_grad = self.params_grad assert type(params_grad) is OrderedDict for ind, k in enumerate(self.params.keys()): updater(index=ind, grad=params_grad[k], weight=self.params[k]) def update_acc_grad(self): if self.acc_grad is None: self.acc_grad = OrderedDict([(n, nd.zeros(v.shape, ctx=self.ctx)) for n, v in self.params_grad.items()]) for k, v in self.acc_grad.items(): v[:] = v + self.params_grad[k] def reset_acc_grad(self): for v in self.acc_grad.values(): v[:] = 0 def copy(self, name=None, ctx=None): if ctx is None: ctx = self.ctx if name is None: name = self.name + '-copy-' + str(ctx) return Base(data_shapes=self.data_shapes, sym_gen=self.sym_gen, default_bucket_kwargs=dict(self.curr_bucket_key), params=self.params, aux_states=self.aux_states, ctx=ctx, name=name) def copy_params_to(self, dst): for k, v in self.params.items(): dst.params[k][:] = v # TODO `wait_to_read()` here seems unnecessary, remove it in the future! dst.params[k].wait_to_read() @property def total_param_num(self): return sum(v.size for v in self.params.values()) def print_stat(self): logging.info("Name: %s" % self.name) assert self.params is not None, "Fatal Error!" logging.info("Params: ") for k, v in self.params.items(): logging.info(" %s: %s" % (k, v.shape)) if self.aux_states is None or 0 == len(self.aux_states): logging.info("Aux States: None") else: logging.info("Aux States: " + ' '.join( ["%s:%s" % (str(k), str(v.shape)) for k, v in self.aux_states.items()])) logging.info("Total Parameter Num: " + str(self.total_param_num))