class ParallModule(BaseModule):
def __init__(self, symbol, data_names, label_names,
logger=logging, context=ctx.cpu(), work_load_list=None,
asymbol = None,
args = None):
super(ParallModule, self).__init__(logger=logger)
self._symbol = symbol
self._asymbol = asymbol
self._data_names = data_names
self._label_names = label_names
self._context = context
self._work_load_list = work_load_list
self._num_classes = config.num_classes
self._batch_size = args.batch_size
self._verbose = args.verbose
self._emb_size = config.emb_size
self._local_class_start = args.local_class_start
self._iter = 0
self._curr_module = None
self._num_workers = config.num_workers
self._num_ctx = len(self._context)
self._ctx_num_classes = args.ctx_num_classes
self._nd_cache = {}
self._ctx_cpu = mx.cpu()
self._ctx_single_gpu = self._context[-1]
self._fixed_param_names = None
self._curr_module = Module(self._symbol, self._data_names, self._label_names, logger=self.logger,
context=self._context, work_load_list=self._work_load_list,
fixed_param_names=self._fixed_param_names)
self._arcface_modules = []
self._ctx_class_start = []
for i in range(len(self._context)):
args._ctxid = i
_module = Module(self._asymbol(args), self._data_names, self._label_names, logger=self.logger,
context=mx.gpu(i), work_load_list=self._work_load_list,
fixed_param_names=self._fixed_param_names)
self._arcface_modules.append(_module)
_c = args.local_class_start + i*args.ctx_num_classes
self._ctx_class_start.append(_c)
self._usekv = False
if self._usekv:
self._distkv = mx.kvstore.create('dist_sync')
self._kvinit = {}
def _reset_bind(self):
self.binded = False
self._curr_module = None
@property
def data_names(self):
return self._data_names
@property
def output_names(self):
return self._symbol.list_outputs()
@property
def data_shapes(self):
assert self.binded
return self._curr_module.data_shapes
@property
def label_shapes(self):
assert self.binded
return self._curr_module.label_shapes
@property
def output_shapes(self):
assert self.binded
return self._curr_module.output_shapes
def get_export_params(self):
assert self.binded and self.params_initialized
_g, _x = self._curr_module.get_params()
g = _g.copy()
x = _x.copy()
return g, x
def get_params(self):
assert self.binded and self.params_initialized
_g, _x = self._curr_module.get_params()
g = _g.copy()
x = _x.copy()
for _module in self._arcface_modules:
_g, _x = _module.get_params()
ag = _g.copy()
ax = _x.copy()
g.update(ag)
x.update(ax)
return g, x
def set_params(self, arg_params, aux_params, allow_missing=False, force_init=True,
allow_extra=False):
g = arg_params
x = aux_params
#ag = {}
#ax = {}
rk = []
for k in g:
v = g[k]
if k.startswith('fc7'):
p1 = k.find('_')
p2 = k.rfind('_')
_ctxid = int(k[p1+1:p2])
self._arcface_modules[_ctxid].set_params({k:v}, {})
rk.append(k)
for k in rk:
del g[k]
self._curr_module.set_params(g, x)
#self._arcface_module.set_params(ag, ax)
def init_params(self, initializer=Uniform(0.01), arg_params=None, aux_params=None,
allow_missing=False, force_init=False, allow_extra=False):
if self.params_initialized and not force_init:
return
assert self.binded, 'call bind before initializing the parameters'
#TODO init the same weights with all work nodes
self._curr_module.init_params(initializer=initializer, arg_params=None,
aux_params=None, allow_missing=allow_missing,
force_init=force_init, allow_extra=allow_extra)
for _module in self._arcface_modules:
#_initializer = initializer
_initializer = mx.init.Normal(0.01)
_module.init_params(initializer=_initializer, arg_params=None,
aux_params=None, allow_missing=allow_missing,
force_init=force_init, allow_extra=allow_extra)
self.params_initialized = True
def bind(self, data_shapes, label_shapes=None, for_training=True,
inputs_need_grad=False, force_rebind=False, shared_module=None):
print('in_bind', self.params_initialized, data_shapes, label_shapes)
if self.params_initialized:
arg_params, aux_params = self.get_params()
# force rebinding is typically used when one want to switch from
# training to prediction phase.
if force_rebind:
self._reset_bind()
if self.binded:
self.logger.warning('Already binded, ignoring bind()')
return
assert shared_module is None, 'shared_module for MutableModule is not supported'
self.for_training = for_training
self.inputs_need_grad = inputs_need_grad
self.binded = True
self._curr_module.bind(data_shapes, label_shapes, for_training, inputs_need_grad,
force_rebind=False, shared_module=None)
_data_shape = data_shapes[0][1]
print('_data_shape', _data_shape, label_shapes)
for _module in self._arcface_modules:
_module.bind([('data', (_data_shape[0]*self._num_workers, self._emb_size))], [('softmax_label', (_data_shape[0]*self._num_workers,))], for_training, True,
force_rebind=False, shared_module=None)
if self.params_initialized:
self.set_params(arg_params, aux_params)
def init_optimizer(self, kvstore='local', optimizer='sgd',
optimizer_params=(('learning_rate', 0.01),), force_init=False):
assert self.binded and self.params_initialized
if self.optimizer_initialized and not force_init:
self.logger.warning('optimizer already initialized, ignoring.')
return
self._curr_module.init_optimizer(kvstore, optimizer, optimizer_params,
force_init=force_init)
for _module in self._arcface_modules:
_module.init_optimizer(kvstore, optimizer, optimizer_params,
force_init=force_init)
self.optimizer_initialized = True
def kv_push(self, key, value):
#if value.context!=mx.cpu():
# value = value.as_in_context(mx.cpu())
if not key in self._kvinit:
self._distkv.init(key, nd.zeros_like(value))
self._kvinit[key] = 1
self._distkv.push(key, value)
#get fc1 and partial fc7
def forward(self, data_batch, is_train=None):
#g,x = self.get_params()
#print('{fc7_weight[0][0]}', self._iter, g['fc7_0_weight'].asnumpy()[0][0])
#print('{pre_fc1_weight[0][0]}', self._iter, g['pre_fc1_weight'].asnumpy()[0][0])
assert self.binded and self.params_initialized
self._curr_module.forward(data_batch, is_train=is_train)
if is_train:
self._iter+=1
fc1, label = self._curr_module.get_outputs(merge_multi_context=True)
global_fc1 = fc1
self.global_label = label.as_in_context(self._ctx_cpu)
for i, _module in enumerate(self._arcface_modules):
_label = self.global_label - self._ctx_class_start[i]
db_global_fc1 = io.DataBatch([global_fc1], [_label])
_module.forward(db_global_fc1) #fc7 with margin
#print('forward end')
def get_ndarray(self, context, name, shape):
key = "%s_%s"%(name, context)
#print(key)
if not key in self._nd_cache:
v = nd.zeros( shape=shape, ctx = context)
self._nd_cache[key] = v
else:
v = self._nd_cache[key]
return v
def get_ndarray2(self, context, name, arr):
key = "%s_%s"%(name, context)
#print(key)
if not key in self._nd_cache:
v = nd.zeros( shape=arr.shape, ctx = context)
self._nd_cache[key] = v
else:
v = self._nd_cache[key]
arr.copyto(v)
return v
def backward(self, out_grads=None):
#print('in backward')
assert self.binded and self.params_initialized
#tmp_ctx = self._ctx_cpu
tmp_ctx = self._ctx_single_gpu
fc7_outs = []
ctx_fc7_max = self.get_ndarray(tmp_ctx, 'ctx_fc7_max', (self._batch_size, len(self._context)))
#local_fc7_max = nd.zeros( (self.global_label.shape[0],1), ctx=mx.cpu())
arcface_module_outputs = []
for i, _module in enumerate(self._arcface_modules):
#_fc7 = _module.get_outputs(merge_multi_context=True)[0]
out = _module.get_outputs(merge_multi_context=True)
#print(out[0].shape)
#print(out[1].shape)
arcface_module_outputs.append(out)
_fc7 = out[0]
fc7_outs.append(_fc7)
_fc7_max = nd.max(_fc7, axis=1).as_in_context(tmp_ctx)
ctx_fc7_max[:,i] = _fc7_max
local_fc7_max = self.get_ndarray(tmp_ctx, 'local_fc7_max', (self._batch_size, 1))
nd.max(ctx_fc7_max, axis=1, keepdims=True, out=local_fc7_max)
global_fc7_max = local_fc7_max
#local_fc7_sum = None
local_fc7_sum = self.get_ndarray(tmp_ctx, 'local_fc7_sum', (self._batch_size,1))
local_fc7_sum[:,:] = 0.0
for i, _module in enumerate(self._arcface_modules):
_max = self.get_ndarray2(fc7_outs[i].context, 'fc7_max', global_fc7_max)
fc7_outs[i] = nd.broadcast_sub(fc7_outs[i], _max)
fc7_outs[i] = nd.exp(fc7_outs[i])
_sum = nd.sum(fc7_outs[i], axis=1, keepdims=True).as_in_context(tmp_ctx)
local_fc7_sum += _sum
global_fc7_sum = local_fc7_sum
if self._iter%self._verbose==0:
#_ctx = self._context[-1]
_ctx = self._ctx_cpu
_probs = []
for i, _module in enumerate(self._arcface_modules):
_prob = self.get_ndarray2(_ctx, '_fc7_prob_%d'%i, fc7_outs[i])
_probs.append(_prob)
fc7_prob = self.get_ndarray(_ctx, 'test_fc7_prob', (self._batch_size, self._ctx_num_classes*len(self._context)))
nd.concat(*_probs, dim=1, out=fc7_prob)
fc7_pred = nd.argmax(fc7_prob, axis=1)
local_label = self.global_label - self._local_class_start
#local_label = self.get_ndarray2(_ctx, 'test_label', local_label)
_pred = nd.equal(fc7_pred, local_label)
print('{fc7_acc}', self._iter, nd.mean(_pred).asnumpy()[0])
#local_fc1_grad = []
#fc1_grad_ctx = self._ctx_cpu
fc1_grad_ctx = self._ctx_single_gpu
local_fc1_grad = self.get_ndarray(fc1_grad_ctx, 'local_fc1_grad', (self._batch_size,self._emb_size))
local_fc1_grad[:,:] = 0.0
total_eloss = []
celoss_verbose = 1000
if self._iter%celoss_verbose==0:
fc7_celoss = self.get_ndarray(tmp_ctx, 'test_fc7_celoss', (self._batch_size,))
fc7_celoss[:] = 0.0
for i, _module in enumerate(self._arcface_modules):
_sum = self.get_ndarray2(fc7_outs[i].context, 'fc7_sum', global_fc7_sum)
fc7_outs[i] = nd.broadcast_div(fc7_outs[i], _sum)
a = i*self._ctx_num_classes
b = (i+1)*self._ctx_num_classes
_label = self.global_label - self._ctx_class_start[i]
_label = self.get_ndarray2(fc7_outs[i].context, 'label', _label)
onehot_label = self.get_ndarray(fc7_outs[i].context, 'label_onehot', (self._batch_size, self._ctx_num_classes))
nd.one_hot(_label, depth=self._ctx_num_classes, on_value = 1.0, off_value = 0.0, out=onehot_label)
#print(fc7_outs[i].shape, onehot_label.shape)
if self._iter%celoss_verbose==0:
_ce_loss = fc7_outs[i] * onehot_label
_ce_loss = nd.sum(_ce_loss, axis=1)
fc7_celoss += _ce_loss.as_in_context(tmp_ctx)
fc7_outs[i] -= onehot_label
out = arcface_module_outputs[i]
out_grads = [fc7_outs[i]]
for j in range(1, len(out)):
eloss = out[j]
#print('eloss%d:'%j, eloss.shape)
#print(out_grads[0].shape)
#egrad_shape = (out_grads[0].shape[0], eloss.shape[0])
egrad_shape = eloss.shape
egrad = self.get_ndarray(fc7_outs[i].context, 'egrad%d'%j, egrad_shape)
#egrad[:][:] = 1.0/egrad_shape[0]
egrad[:][:] = 1.0
out_grads.append(egrad)
if self._iter%self._verbose==0:
total_eloss.append(np.mean(eloss.asnumpy()))
_module.backward(out_grads = out_grads)
#ctx_fc1_grad = _module.get_input_grads()[0].as_in_context(mx.cpu())
ctx_fc1_grad = self.get_ndarray2(fc1_grad_ctx, 'ctx_fc1_grad_%d'%i, _module.get_input_grads()[0])
local_fc1_grad += ctx_fc1_grad
if self._iter%self._verbose==0 and len(total_eloss)>0:
print('{eloss}', self._iter, np.mean(total_eloss))
#if self._iter%self._verbose==0:
if self._iter%celoss_verbose==0:
ce_loss = nd.log(fc7_celoss) * -1.0
ce_loss = nd.mean(ce_loss)
print('CELOSS,%d,%f'% (self._iter, ce_loss.asnumpy()))
global_fc1_grad = local_fc1_grad
self._curr_module.backward(out_grads = [global_fc1_grad])
def update(self):
assert self.binded and self.params_initialized and self.optimizer_initialized
self._curr_module.update()
for i, _module in enumerate(self._arcface_modules):
_module.update()
mx.nd.waitall()
def get_outputs(self, merge_multi_context=True):
assert self.binded and self.params_initialized
return self._curr_module.get_outputs(merge_multi_context=merge_multi_context)
#return self._arcface_module.get_outputs(merge_multi_context=merge_multi_context)
def get_input_grads(self, merge_multi_context=True):
assert self.binded and self.params_initialized and self.inputs_need_grad
return self._curr_module.get_input_grads(merge_multi_context=merge_multi_context)
def update_metric(self, eval_metric, labels):
assert self.binded and self.params_initialized
#self._curr_module.update_metric(eval_metric, labels)
#label = labels[0]
#print(label.shape)
#self._arcface_module.update_metric(eval_metric, labels)
def install_monitor(self, mon):
""" Install monitor on all executors """
assert self.binded
self._curr_module.install_monitor(mon)
def forward_backward(self, data_batch):
"""A convenient function that calls both ``forward`` and ``backward``."""
self.forward(data_batch, is_train=True) # get fc1 and partial fc7
self.backward()
def fit(self, train_data, eval_data=None, eval_metric='acc',
epoch_end_callback=None, batch_end_callback=None, kvstore='local',
optimizer='sgd', optimizer_params=(('learning_rate', 0.01),),
eval_end_callback=None,
eval_batch_end_callback=None, initializer=Uniform(0.01),
arg_params=None, aux_params=None, allow_missing=False,
force_rebind=False, force_init=False, begin_epoch=0, num_epoch=None,
validation_metric=None, monitor=None, sparse_row_id_fn=None):
"""Trains the module parameters.
Checkout `Module Tutorial <http://mxnet.io/tutorials/basic/module.html>`_ to see
a end-to-end use-case.
Parameters
----------
train_data : DataIter
Train DataIter.
eval_data : DataIter
If not ``None``, will be used as validation set and the performance
after each epoch will be evaluated.
eval_metric : str or EvalMetric
Defaults to 'accuracy'. The performance measure used to display during training.
Other possible predefined metrics are:
'ce' (CrossEntropy), 'f1', 'mae', 'mse', 'rmse', 'top_k_accuracy'.
epoch_end_callback : function or list of functions
Each callback will be called with the current `epoch`, `symbol`, `arg_params`
and `aux_params`.
batch_end_callback : function or list of function
Each callback will be called with a `BatchEndParam`.
kvstore : str or KVStore
Defaults to 'local'.
optimizer : str or Optimizer
Defaults to 'sgd'.
optimizer_params : dict
Defaults to ``(('learning_rate', 0.01),)``. The parameters for
the optimizer constructor.
The default value is not a dict, just to avoid pylint warning on dangerous
default values.
eval_end_callback : function or list of function
These will be called at the end of each full evaluation, with the metrics over
the entire evaluation set.
eval_batch_end_callback : function or list of function
These will be called at the end of each mini-batch during evaluation.
initializer : Initializer
The initializer is called to initialize the module parameters when they are
not already initialized.
arg_params : dict
Defaults to ``None``, if not ``None``, should be existing parameters from a trained
model or loaded from a checkpoint (previously saved model). In this case,
the value here will be used to initialize the module parameters, unless they
are already initialized by the user via a call to `init_params` or `fit`.
`arg_params` has a higher priority than `initializer`.
aux_params : dict
Defaults to ``None``. Similar to `arg_params`, except for auxiliary states.
allow_missing : bool
Defaults to ``False``. Indicates whether to allow missing parameters when `arg_params`
and `aux_params` are not ``None``. If this is ``True``, then the missing parameters
will be initialized via the `initializer`.
force_rebind : bool
Defaults to ``False``. Whether to force rebinding the executors if already bound.
force_init : bool
Defaults to ``False``. Indicates whether to force initialization even if the
parameters are already initialized.
begin_epoch : int
Defaults to 0. Indicates the starting epoch. Usually, if resumed from a
checkpoint saved at a previous training phase at epoch N, then this value should be
N+1.
num_epoch : int
Number of epochs for training.
sparse_row_id_fn : A callback function
The function takes `data_batch` as an input and returns a dict of
str -> NDArray. The resulting dict is used for pulling row_sparse
parameters from the kvstore, where the str key is the name of the param,
and the value is the row id of the param to pull.
Examples
--------
>>> # An example of using fit for training.
>>> # Assume training dataIter and validation dataIter are ready
>>> # Assume loading a previously checkpointed model
>>> sym, arg_params, aux_params = mx.model.load_checkpoint(model_prefix, 3)
>>> mod.fit(train_data=train_dataiter, eval_data=val_dataiter, optimizer='sgd',
... optimizer_params={'learning_rate':0.01, 'momentum': 0.9},
... arg_params=arg_params, aux_params=aux_params,
... eval_metric='acc', num_epoch=10, begin_epoch=3)
"""
assert num_epoch is not None, 'please specify number of epochs'
assert arg_params is None and aux_params is None
self.bind(data_shapes=train_data.provide_data, label_shapes=train_data.provide_label,
for_training=True, force_rebind=force_rebind)
if monitor is not None:
self.install_monitor(monitor)
self.init_params(initializer=initializer, arg_params=arg_params, aux_params=aux_params,
allow_missing=allow_missing, force_init=force_init)
self.init_optimizer(kvstore=kvstore, optimizer=optimizer,
optimizer_params=optimizer_params)
if validation_metric is None:
validation_metric = eval_metric
if not isinstance(eval_metric, metric.EvalMetric):
eval_metric = metric.create(eval_metric)
epoch_eval_metric = copy.deepcopy(eval_metric)
################################################################################
# training loop
################################################################################
for epoch in range(begin_epoch, num_epoch):
tic = time.time()
eval_metric.reset()
epoch_eval_metric.reset()
nbatch = 0
data_iter = iter(train_data)
end_of_batch = False
next_data_batch = next(data_iter)
while not end_of_batch:
data_batch = next_data_batch
if monitor is not None:
monitor.tic()
self.forward_backward(data_batch)
self.update()
assert not isinstance(data_batch, list)
#if isinstance(data_batch, list):
# #print('XXX')
# self.update_metric(eval_metric,
# [db.label for db in data_batch],
# pre_sliced=True)
# self.update_metric(epoch_eval_metric,
# [db.label for db in data_batch],
# pre_sliced=True)
#else:
# #print('before update metric')
# self.update_metric(eval_metric, data_batch.label)
# self.update_metric(epoch_eval_metric, data_batch.label)
#labels = data_batch.label
#labels = [self.global_label]
#self.update_metric(eval_metric, labels)
#self.update_metric(epoch_eval_metric, labels)
try:
# pre fetch next batch
next_data_batch = next(data_iter)
self.prepare(next_data_batch, sparse_row_id_fn=sparse_row_id_fn)
except StopIteration:
end_of_batch = True
if monitor is not None:
monitor.toc_print()
#if end_of_batch:
# eval_name_vals = epoch_eval_metric.get_name_value()
if batch_end_callback is not None:
batch_end_params = BatchEndParam(epoch=epoch, nbatch=nbatch,
eval_metric=None,
locals=locals())
batch_end_callback(batch_end_params)
#for callback in _as_list(batch_end_callback):
# callback(batch_end_params)
nbatch += 1
# one epoch of training is finished
#for name, val in eval_name_vals:
# self.logger.info('Epoch[%d] Train-%s=%f', epoch, name, val)
toc = time.time()
self.logger.info('Epoch[%d] Time cost=%.3f', epoch, (toc-tic))
# sync aux params across devices
arg_params, aux_params = self.get_params()
self.set_params(arg_params, aux_params)
# end of 1 epoch, reset the data-iter for another epoch
train_data.reset()
