def train_net(args):
ctx = []
cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
if len(cvd) > 0:
for i in xrange(len(cvd.split(','))):
ctx.append(mx.gpu(i))
if len(ctx) == 0:
ctx = [mx.cpu()]
print('use cpu')
else:
print('gpu num:', len(ctx))
prefix = args.prefix
prefix_dir = os.path.dirname(prefix)
if not os.path.exists(prefix_dir):
os.makedirs(prefix_dir)
end_epoch = args.end_epoch
args.ctx_num = len(ctx)
args.num_layers = int(args.network[1:])
print('num_layers', args.num_layers)
if args.per_batch_size == 0:
args.per_batch_size = 128
args.batch_size = args.per_batch_size * args.ctx_num
args.image_channel = 3
data_dir = args.data_dir
if args.task == 'gender':
data_dir = args.gender_data_dir
elif args.task == 'age':
data_dir = args.age_data_dir
print('data dir', data_dir)
path_imgrec = None
path_imglist = None
prop = face_image.load_property(data_dir)
args.num_classes = prop.num_classes
image_size = prop.image_size
args.image_h = image_size[0]
args.image_w = image_size[1]
print('image_size', image_size)
assert (args.num_classes > 0)
print('num_classes', args.num_classes)
path_imgrec = os.path.join(data_dir, "train.rec")
print('Called with argument:', args)
data_shape = (args.image_channel, image_size[0], image_size[1])
mean = None
begin_epoch = 0
net = get_model()
#if args.task=='':
# test_net = get_model_test(net)
#print(net.__class__)
#net = net0[0]
if args.network[0] == 'r' or args.network[0] == 'y':
initializer = mx.init.Xavier(rnd_type='gaussian',
factor_type="out",
magnitude=2) #resnet style
elif args.network[0] == 'i' or args.network[0] == 'x':
initializer = mx.init.Xavier(rnd_type='gaussian',
factor_type="in",
magnitude=2) #inception
else:
initializer = mx.init.Xavier(rnd_type='uniform',
factor_type="in",
magnitude=2)
net.hybridize()
if args.mode == 'gluon':
if len(args.pretrained) == 0:
pass
else:
net.load_params(args.pretrained,
allow_missing=True,
ignore_extra=True)
net.initialize(initializer)
net.collect_params().reset_ctx(ctx)
val_iter = None
if args.task == '':
train_iter = FaceImageIter(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=path_imgrec,
shuffle=True,
rand_mirror=args.rand_mirror,
mean=mean,
cutoff=args.cutoff,
)
else:
train_iter = FaceImageIterAge(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=path_imgrec,
task=args.task,
shuffle=True,
rand_mirror=args.rand_mirror,
mean=mean,
cutoff=args.cutoff,
)
if args.task == 'age':
metric = CompositeEvalMetric([MAEMetric(), CUMMetric()])
elif args.task == 'gender':
metric = CompositeEvalMetric([AccMetric()])
else:
metric = CompositeEvalMetric([AccMetric()])
ver_list = []
ver_name_list = []
if args.task == '':
for name in args.eval.split(','):
path = os.path.join(data_dir, name + ".bin")
if os.path.exists(path):
data_set = verification.load_bin(path, image_size)
ver_list.append(data_set)
ver_name_list.append(name)
print('ver', name)
def ver_test(nbatch):
results = []
for i in xrange(len(ver_list)):
acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(
ver_list[i], net, ctx, batch_size=args.batch_size)
print('[%s][%d]XNorm: %f' % (ver_name_list[i], nbatch, xnorm))
#print('[%s][%d]Accuracy: %1.5f+-%1.5f' % (ver_name_list[i], nbatch, acc1, std1))
print('[%s][%d]Accuracy-Flip: %1.5f+-%1.5f' %
(ver_name_list[i], nbatch, acc2, std2))
results.append(acc2)
return results
def val_test(nbatch=0):
acc = 0.0
#if args.task=='age':
if len(args.age_data_dir) > 0:
val_iter = FaceImageIterAge(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=os.path.join(args.age_data_dir, 'val.rec'),
task=args.task,
shuffle=False,
rand_mirror=False,
mean=mean,
)
_metric = MAEMetric()
val_metric = mx.metric.create(_metric)
val_metric.reset()
_metric2 = CUMMetric()
val_metric2 = mx.metric.create(_metric2)
val_metric2.reset()
val_iter.reset()
for batch in val_iter:
data = gluon.utils.split_and_load(batch.data[0],
ctx_list=ctx,
batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0],
ctx_list=ctx,
batch_axis=0)
outputs = []
for x in data:
outputs.append(net(x)[2])
val_metric.update(label, outputs)
val_metric2.update(label, outputs)
_value = val_metric.get_name_value()[0][1]
print('[%d][VMAE]: %f' % (nbatch, _value))
_value = val_metric2.get_name_value()[0][1]
if args.task == 'age':
acc = _value
print('[%d][VCUM]: %f' % (nbatch, _value))
if len(args.gender_data_dir) > 0:
val_iter = FaceImageIterAge(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=os.path.join(args.gender_data_dir, 'val.rec'),
task=args.task,
shuffle=False,
rand_mirror=False,
mean=mean,
)
_metric = AccMetric()
val_metric = mx.metric.create(_metric)
val_metric.reset()
val_iter.reset()
for batch in val_iter:
data = gluon.utils.split_and_load(batch.data[0],
ctx_list=ctx,
batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0],
ctx_list=ctx,
batch_axis=0)
outputs = []
for x in data:
outputs.append(net(x)[1])
val_metric.update(label, outputs)
_value = val_metric.get_name_value()[0][1]
if args.task == 'gender':
acc = _value
print('[%d][VACC]: %f' % (nbatch, _value))
return acc
total_time = 0
num_epochs = 0
best_acc = [0]
highest_acc = [0.0, 0.0] #lfw and target
global_step = [0]
save_step = [0]
if len(args.lr_steps) == 0:
lr_steps = [100000, 140000, 160000]
p = 512.0 / args.batch_size
for l in xrange(len(lr_steps)):
lr_steps[l] = int(lr_steps[l] * p)
else:
lr_steps = [int(x) for x in args.lr_steps.split(',')]
print('lr_steps', lr_steps)
kv = mx.kv.create('device')
#kv = mx.kv.create('local')
#_rescale = 1.0/args.ctx_num
#opt = optimizer.SGD(learning_rate=args.lr, momentum=args.mom, wd=args.wd, rescale_grad=_rescale)
#opt = optimizer.SGD(learning_rate=args.lr, momentum=args.mom, wd=args.wd)
if args.mode == 'gluon':
trainer = gluon.Trainer(net.collect_params(),
'sgd', {
'learning_rate': args.lr,
'wd': args.wd,
'momentum': args.mom,
'multi_precision': True
},
kvstore=kv)
else:
_rescale = 1.0 / args.ctx_num
opt = optimizer.SGD(learning_rate=args.lr,
momentum=args.mom,
wd=args.wd,
rescale_grad=_rescale)
_cb = mx.callback.Speedometer(args.batch_size, 20)
arg_params = None
aux_params = None
data = mx.sym.var('data')
label = mx.sym.var('softmax_label')
if args.margin_a > 0.0:
fc7 = net(data, label)
else:
fc7 = net(data)
#sym = mx.symbol.SoftmaxOutput(data=fc7, label = label, name='softmax', normalization='valid')
ceop = gluon.loss.SoftmaxCrossEntropyLoss()
loss = ceop(fc7, label)
#loss = loss/args.per_batch_size
loss = mx.sym.mean(loss)
sym = mx.sym.Group([
mx.symbol.BlockGrad(fc7),
mx.symbol.MakeLoss(loss, name='softmax')
])
def _batch_callback():
mbatch = global_step[0]
global_step[0] += 1
for _lr in lr_steps:
if mbatch == _lr:
args.lr *= 0.1
if args.mode == 'gluon':
trainer.set_learning_rate(args.lr)
else:
opt.lr = args.lr
print('lr change to', args.lr)
break
#_cb(param)
if mbatch % 1000 == 0:
print('lr-batch-epoch:', args.lr, mbatch)
if mbatch > 0 and mbatch % args.verbose == 0:
save_step[0] += 1
msave = save_step[0]
do_save = False
is_highest = False
if args.task == 'age' or args.task == 'gender':
acc = val_test(mbatch)
if acc >= highest_acc[-1]:
highest_acc[-1] = acc
is_highest = True
do_save = True
else:
acc_list = ver_test(mbatch)
if len(acc_list) > 0:
lfw_score = acc_list[0]
if lfw_score > highest_acc[0]:
highest_acc[0] = lfw_score
if lfw_score >= 0.998:
do_save = True
if acc_list[-1] >= highest_acc[-1]:
highest_acc[-1] = acc_list[-1]
if lfw_score >= 0.99:
do_save = True
is_highest = True
if args.ckpt == 0:
do_save = False
elif args.ckpt > 1:
do_save = True
if do_save:
print('saving', msave)
#print('saving gluon params')
fname = os.path.join(args.prefix, 'model-gluon.params')
net.save_params(fname)
fname = os.path.join(args.prefix, 'model')
net.export(fname, msave)
#arg, aux = model.get_params()
#mx.model.save_checkpoint(prefix, msave, model.symbol, arg, aux)
print('[%d]Accuracy-Highest: %1.5f' % (mbatch, highest_acc[-1]))
if args.max_steps > 0 and mbatch > args.max_steps:
sys.exit(0)
def _batch_callback_sym(param):
_cb(param)
_batch_callback()
if args.mode != 'gluon':
model = mx.mod.Module(
context=ctx,
symbol=sym,
)
model.fit(train_iter,
begin_epoch=0,
num_epoch=args.end_epoch,
eval_data=None,
eval_metric=metric,
kvstore='device',
optimizer=opt,
initializer=initializer,
arg_params=arg_params,
aux_params=aux_params,
allow_missing=True,
batch_end_callback=_batch_callback_sym,
epoch_end_callback=None)
else:
loss_weight = 1.0
if args.task == 'age':
loss_weight = 1.0 / AGE
#loss = gluon.loss.SoftmaxCrossEntropyLoss(weight = loss_weight)
loss = nd.SoftmaxOutput
#loss = gluon.loss.SoftmaxCrossEntropyLoss()
while True:
#trainer = update_learning_rate(opt.lr, trainer, epoch, opt.lr_factor, lr_steps)
tic = time.time()
train_iter.reset()
metric.reset()
btic = time.time()
for i, batch in enumerate(train_iter):
_batch_callback()
#data = gluon.utils.split_and_load(batch.data[0].astype(opt.dtype), ctx_list=ctx, batch_axis=0)
#label = gluon.utils.split_and_load(batch.label[0].astype(opt.dtype), ctx_list=ctx, batch_axis=0)
data = gluon.utils.split_and_load(batch.data[0],
ctx_list=ctx,
batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0],
ctx_list=ctx,
batch_axis=0)
outputs = []
Ls = []
with ag.record():
for x, y in zip(data, label):
#print(y.asnumpy())
if args.task == '':
if args.margin_a > 0.0:
z = net(x, y)
else:
z = net(x)
#print(z[0].shape, z[1].shape)
else:
z = net(x)
if args.task == 'gender':
L = loss(z[1], y)
#L = L/args.per_batch_size
Ls.append(L)
outputs.append(z[1])
elif args.task == 'age':
for k in xrange(AGE):
_z = nd.slice_axis(z[2],
axis=1,
begin=k * 2,
end=k * 2 + 2)
_y = nd.slice_axis(y,
axis=1,
begin=k,
end=k + 1)
_y = nd.flatten(_y)
L = loss(_z, _y)
#L = L/args.per_batch_size
#L /= AGE
Ls.append(L)
outputs.append(z[2])
else:
L = loss(z, y)
#L = L/args.per_batch_size
Ls.append(L)
outputs.append(z)
# store the loss and do backward after we have done forward
# on all GPUs for better speed on multiple GPUs.
ag.backward(Ls)
#trainer.step(batch.data[0].shape[0], ignore_stale_grad=True)
#trainer.step(args.ctx_num)
n = batch.data[0].shape[0]
#print(n,n)
trainer.step(n)
metric.update(label, outputs)
if i > 0 and i % 20 == 0:
name, acc = metric.get()
if len(name) == 2:
logger.info(
'Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f, %s=%f'
% (num_epochs, i, args.batch_size /
(time.time() - btic), name[0], acc[0], name[1],
acc[1]))
else:
logger.info(
'Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f'
% (num_epochs, i, args.batch_size /
(time.time() - btic), name[0], acc[0]))
#metric.reset()
btic = time.time()
epoch_time = time.time() - tic
# First epoch will usually be much slower than the subsequent epics,
# so don't factor into the average
if num_epochs > 0:
total_time = total_time + epoch_time
#name, acc = metric.get()
#logger.info('[Epoch %d] training: %s=%f, %s=%f'%(num_epochs, name[0], acc[0], name[1], acc[1]))
logger.info('[Epoch %d] time cost: %f' % (num_epochs, epoch_time))
num_epochs = num_epochs + 1
#name, val_acc = test(ctx, val_data)
#logger.info('[Epoch %d] validation: %s=%f, %s=%f'%(epoch, name[0], val_acc[0], name[1], val_acc[1]))
# save model if meet requirements
#save_checkpoint(epoch, val_acc[0], best_acc)
if num_epochs > 1:
print('Average epoch time: {}'.format(
float(total_time) / (num_epochs - 1)))
class MXNetRunner(object):
"""Manages a MXNet model for training."""
def setup_distributed(self, env, config, model_creator, loss_creator=None,
validation_metrics_creator=None, eval_metrics_creator=None):
logging.basicConfig(level=logging.INFO) # This can print log messages to console.
self.logger = logging.getLogger()
assert isinstance(config, dict), "config must be a dict"
for param in ["optimizer", "optimizer_params", "log_interval"]:
assert param in config, param + " must be specified in config"
self.config = config
self.model_creator = model_creator
self.loss_creator = loss_creator
self.validation_metrics_creator = validation_metrics_creator
self.eval_metrics_creator = eval_metrics_creator
self.is_worker = False
env["DMLC_NODE_HOST"] = self.get_node_ip()
if env["DMLC_ROLE"] == "worker":
self.is_worker = True
if self.is_worker:
os.environ.update(env)
self.kv = mx.kv.create("dist_sync")
# Set seed so that the model on each worker is initialized with the same weights.
if "seed" in self.config:
mx.random.seed(self.config["seed"])
self.model = self.model_creator(self.config)
self.loss = self.loss_creator(self.config) if self.loss_creator else None
self.eval_metrics = self.eval_metrics_creator(self.config) \
if self.eval_metrics_creator else None
from mxnet.metric import CompositeEvalMetric
if isinstance(self.eval_metrics, list):
self.eval_metrics = CompositeEvalMetric(self.eval_metrics)
self.val_metrics = self.validation_metrics_creator(self.config) \
if self.validation_metrics_creator else None
if isinstance(self.val_metrics, list):
self.val_metrics = CompositeEvalMetric(self.val_metrics)
# For BaseModule, use symbolic API. Otherwise, use imperative API.
# TODO: change Gluon Trainer to Estimator API?
if not isinstance(self.model, mx.module.BaseModule):
assert self.loss, "Loss not defined for gluon model, please specify loss_creator"
self.trainer = gluon.Trainer(self.model.collect_params(), self.config["optimizer"],
optimizer_params=self.config["optimizer_params"],
kvstore=self.kv)
else: # Trainer is not needed for symbolic API.
self.trainer = None
else: # server
# Need to use the environment on each raylet process for the correct python environment.
# TODO: Need to kill this process manually?
modified_env = os.environ.copy()
modified_env.update(env)
# For servers, just import mxnet and no need to do anything else.
subprocess.Popen("python -c 'import mxnet'", shell=True, env=modified_env)
def train(self, train_data, epochs=1, batch_size=32,
validation_data=None, train_resize_batch_num=None):
"""Train the model and update the model parameters."""
stats = dict()
if self.is_worker:
from zoo.orca.data.shard import RayPartition
if isinstance(train_data, RayPartition):
from zoo.orca.data.utils import ray_partition_get_data_label
data, label = ray_partition_get_data_label(train_data.get_data(),
allow_tuple=False,
allow_list=False)
train_data_iter = mx.io.NDArrayIter(data=data, label=label,
batch_size=batch_size, shuffle=True)
if train_resize_batch_num is not None:
train_data_iter = mx.io.ResizeIter(train_data_iter, train_resize_batch_num)
if validation_data:
data_val, label_val = ray_partition_get_data_label(validation_data.get_data(),
allow_tuple=False,
allow_list=False)
val_data_iter = mx.io.NDArrayIter(data=data_val, label=label_val,
batch_size=batch_size, shuffle=True)
else:
val_data_iter = None
else: # data_creator functions; should return Iter or DataLoader
config = self.config
if "batch_size" not in config:
config["batch_size"] = batch_size
train_data_iter = train_data(config, self.kv)
val_data_iter = validation_data(config, self.kv) if validation_data else None
start_time = time.time()
if self.trainer: # Imperative API
def cpu_context(target_data):
if isinstance(target_data, list):
return [cpu_context(d) for d in target_data]
else:
return target_data.as_in_context(mx.cpu())
for epoch in range(epochs):
# DataLoader doesn't need to be reset.
if isinstance(train_data_iter, mx.io.DataIter):
train_data_iter.reset()
if self.eval_metrics:
self.eval_metrics.reset() # metrics will accumulate for one batch.
batch_start_time = time.time()
epoch_start_time = time.time()
for i, batch in enumerate(train_data_iter):
data = cpu_context(batch.data)
label = cpu_context(batch.label)
if not isinstance(data, list):
data = [data]
if not isinstance(label, list):
label = [label]
from mxnet import autograd as ag
with ag.record():
output = self.model(*data) # forward
if not isinstance(output, list):
output = [output]
Ls = self.loss(*output, *label)
ag.backward(Ls)
self.trainer.step(batch_size)
if self.eval_metrics:
self.eval_metrics.update(label, output)
if not (i + 1) % self.config["log_interval"]:
# This would be logged on driver for each worker process.
iteration_log = \
"Epoch[%d] Batch[%d] Speed: %f samples/sec %s=%f" \
% (epoch, i,
batch_size / (time.time() - batch_start_time),
"loss", Ls.asnumpy().mean())
if self.eval_metrics:
names, accs = self.eval_metrics.get()
names, accs = to_list(names), to_list(accs)
for name, acc in zip(names, accs):
iteration_log += " %s=%f" % (name, acc)
self.logger.info(iteration_log)
batch_start_time = time.time()
# Epoch time log.
self.logger.info("[Epoch %d] time cost: %f" %
(epoch, time.time() - epoch_start_time))
# Epoch metrics log on train data.
if self.eval_metrics:
epoch_train_log = "[Epoch %d] training: " % epoch
names, accs = self.eval_metrics.get()
names, accs = to_list(names), to_list(accs)
for name, acc in zip(names, accs):
epoch_train_log += "%s=%f " % (name, acc)
self.logger.info(epoch_train_log)
# Epoch metrics log on validation data if any.
if val_data_iter:
if isinstance(val_data_iter, mx.io.DataIter):
val_data_iter.reset()
self.val_metrics.reset()
for batch in val_data_iter:
data = cpu_context(batch.data)
label = cpu_context(batch.label)
if not isinstance(data, list):
data = [data]
if not isinstance(label, list):
label = [label]
output = self.model(*data)
if not isinstance(output, list):
output = [output]
self.val_metrics.update(label, output)
epoch_val_log = "[Epoch %d] validation: " % epoch
names, accs = self.val_metrics.get()
names, accs = to_list(names), to_list(accs)
for name, acc in zip(names, accs):
epoch_val_log += "%s=%f " % (name, acc)
self.logger.info(epoch_val_log)
# TODO: save checkpoints
if self.eval_metrics:
names, accs = self.eval_metrics.get()
names, accs = to_list(names), to_list(accs)
for name, acc in zip(names, accs):
stats[name] = acc
else: # Symbolic API
# TODO: seems no history (i.e. validation accuracy) returned by fit?
if "init" not in self.config:
from mxnet.initializer import Uniform
self.config["init"] = Uniform(0.01) # This is the default value for MXNet.
if self.eval_metrics is None:
self.eval_metrics = 'acc' # This is the default value for MXNet.
self.model.fit(train_data=train_data_iter,
num_epoch=epochs,
initializer=self.config["init"],
kvstore=self.kv,
optimizer=self.config["optimizer"],
optimizer_params=self.config["optimizer_params"],
eval_data=val_data_iter,
eval_metric=self.eval_metrics,
validation_metric=self.val_metrics,
batch_end_callback=mx.callback.Speedometer(
batch_size, self.config["log_interval"]),
epoch_end_callback=None if "model" not in self.config
else mx.callback.do_checkpoint(self.config["model"]))
epoch_time = time.time() - start_time
stats["epoch_time"] = epoch_time
if isinstance(train_data, RayPartition):
del train_data
if validation_data and isinstance(validation_data, RayPartition):
del validation_data
return stats
def shutdown(self):
"""Attempts to shut down the runner."""
del self.logger
if self.is_worker:
del self.kv
del self.model
del self.trainer
del self.loss
del self.eval_metrics
del self.val_metrics
def get_node_ip(self):
"""Returns the IP address of the current node."""
if "node_ip" not in self.__dict__:
self.node_ip = ray.services.get_node_ip_address()
return self.node_ip
def find_free_port(self):
"""Finds a free port on the current node."""
if "port" not in self.__dict__:
from zoo.orca.learn.mxnet.utils import find_free_port
self.port = find_free_port()
return self.port
def train_net(args):
ctx = []
cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
if len(cvd)>0:
for i in xrange(len(cvd.split(','))):
ctx.append(mx.gpu(i))
if len(ctx)==0:
ctx = [mx.cpu()]
print('use cpu')
else:
print('gpu num:', len(ctx))
prefix = args.prefix
prefix_dir = os.path.dirname(prefix)
if not os.path.exists(prefix_dir):
os.makedirs(prefix_dir)
end_epoch = args.end_epoch
args.ctx_num = len(ctx)
args.num_layers = int(args.network[1:])
print('num_layers', args.num_layers)
if args.per_batch_size==0:
args.per_batch_size = 128
args.batch_size = args.per_batch_size*args.ctx_num
args.image_channel = 3
data_dir = args.data_dir
if args.task=='gender':
data_dir = args.gender_data_dir
elif args.task=='age':
data_dir = args.age_data_dir
print('data dir', data_dir)
path_imgrec = None
path_imglist = None
prop = face_image.load_property(data_dir)
args.num_classes = prop.num_classes
image_size = prop.image_size
args.image_h = image_size[0]
args.image_w = image_size[1]
print('image_size', image_size)
assert(args.num_classes>0)
print('num_classes', args.num_classes)
path_imgrec = os.path.join(data_dir, "train.rec")
print('Called with argument:', args)
data_shape = (args.image_channel,image_size[0],image_size[1])
mean = None
begin_epoch = 0
net = get_model()
#if args.task=='':
# test_net = get_model_test(net)
#print(net.__class__)
#net = net0[0]
if args.network[0]=='r' or args.network[0]=='y':
initializer = mx.init.Xavier(rnd_type='gaussian', factor_type="out", magnitude=2) #resnet style
elif args.network[0]=='i' or args.network[0]=='x':
initializer = mx.init.Xavier(rnd_type='gaussian', factor_type="in", magnitude=2) #inception
else:
initializer = mx.init.Xavier(rnd_type='uniform', factor_type="in", magnitude=2)
net.hybridize()
if args.mode=='gluon':
if len(args.pretrained)==0:
pass
else:
net.load_params(args.pretrained, allow_missing=True, ignore_extra = True)
net.initialize(initializer)
net.collect_params().reset_ctx(ctx)
val_iter = None
if args.task=='':
train_iter = FaceImageIter(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = path_imgrec,
shuffle = True,
rand_mirror = args.rand_mirror,
mean = mean,
cutoff = args.cutoff,
)
else:
train_iter = FaceImageIterAge(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = path_imgrec,
task = args.task,
shuffle = True,
rand_mirror = args.rand_mirror,
mean = mean,
cutoff = args.cutoff,
)
if args.task=='age':
metric = CompositeEvalMetric([MAEMetric(), CUMMetric()])
elif args.task=='gender':
metric = CompositeEvalMetric([AccMetric()])
else:
metric = CompositeEvalMetric([AccMetric()])
ver_list = []
ver_name_list = []
if args.task=='':
for name in args.eval.split(','):
path = os.path.join(data_dir,name+".bin")
if os.path.exists(path):
data_set = verification.load_bin(path, image_size)
ver_list.append(data_set)
ver_name_list.append(name)
print('ver', name)
def ver_test(nbatch):
results = []
for i in xrange(len(ver_list)):
acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(ver_list[i], net, ctx, batch_size = args.batch_size)
print('[%s][%d]XNorm: %f' % (ver_name_list[i], nbatch, xnorm))
#print('[%s][%d]Accuracy: %1.5f+-%1.5f' % (ver_name_list[i], nbatch, acc1, std1))
print('[%s][%d]Accuracy-Flip: %1.5f+-%1.5f' % (ver_name_list[i], nbatch, acc2, std2))
results.append(acc2)
return results
def val_test(nbatch=0):
acc = 0.0
#if args.task=='age':
if len(args.age_data_dir)>0:
val_iter = FaceImageIterAge(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = os.path.join(args.age_data_dir, 'val.rec'),
task = args.task,
shuffle = False,
rand_mirror = False,
mean = mean,
)
_metric = MAEMetric()
val_metric = mx.metric.create(_metric)
val_metric.reset()
_metric2 = CUMMetric()
val_metric2 = mx.metric.create(_metric2)
val_metric2.reset()
val_iter.reset()
for batch in val_iter:
data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
outputs = []
for x in data:
outputs.append(net(x)[2])
val_metric.update(label, outputs)
val_metric2.update(label, outputs)
_value = val_metric.get_name_value()[0][1]
print('[%d][VMAE]: %f'%(nbatch, _value))
_value = val_metric2.get_name_value()[0][1]
if args.task=='age':
acc = _value
print('[%d][VCUM]: %f'%(nbatch, _value))
if len(args.gender_data_dir)>0:
val_iter = FaceImageIterAge(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = os.path.join(args.gender_data_dir, 'val.rec'),
task = args.task,
shuffle = False,
rand_mirror = False,
mean = mean,
)
_metric = AccMetric()
val_metric = mx.metric.create(_metric)
val_metric.reset()
val_iter.reset()
for batch in val_iter:
data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
outputs = []
for x in data:
outputs.append(net(x)[1])
val_metric.update(label, outputs)
_value = val_metric.get_name_value()[0][1]
if args.task=='gender':
acc = _value
print('[%d][VACC]: %f'%(nbatch, _value))
return acc
total_time = 0
num_epochs = 0
best_acc = [0]
highest_acc = [0.0, 0.0] #lfw and target
global_step = [0]
save_step = [0]
if len(args.lr_steps)==0:
lr_steps = [100000, 140000, 160000]
p = 512.0/args.batch_size
for l in xrange(len(lr_steps)):
lr_steps[l] = int(lr_steps[l]*p)
else:
lr_steps = [int(x) for x in args.lr_steps.split(',')]
print('lr_steps', lr_steps)
kv = mx.kv.create('device')
#kv = mx.kv.create('local')
#_rescale = 1.0/args.ctx_num
#opt = optimizer.SGD(learning_rate=args.lr, momentum=args.mom, wd=args.wd, rescale_grad=_rescale)
#opt = optimizer.SGD(learning_rate=args.lr, momentum=args.mom, wd=args.wd)
if args.mode=='gluon':
trainer = gluon.Trainer(net.collect_params(), 'sgd',
{'learning_rate': args.lr, 'wd': args.wd, 'momentum': args.mom, 'multi_precision': True},
kvstore=kv)
else:
_rescale = 1.0/args.ctx_num
opt = optimizer.SGD(learning_rate=args.lr, momentum=args.mom, wd=args.wd, rescale_grad=_rescale)
_cb = mx.callback.Speedometer(args.batch_size, 20)
arg_params = None
aux_params = None
data = mx.sym.var('data')
label = mx.sym.var('softmax_label')
if args.margin_a>0.0:
fc7 = net(data, label)
else:
fc7 = net(data)
#sym = mx.symbol.SoftmaxOutput(data=fc7, label = label, name='softmax', normalization='valid')
ceop = gluon.loss.SoftmaxCrossEntropyLoss()
loss = ceop(fc7, label)
#loss = loss/args.per_batch_size
loss = mx.sym.mean(loss)
sym = mx.sym.Group( [mx.symbol.BlockGrad(fc7), mx.symbol.MakeLoss(loss, name='softmax')] )
def _batch_callback():
mbatch = global_step[0]
global_step[0]+=1
for _lr in lr_steps:
if mbatch==_lr:
args.lr *= 0.1
if args.mode=='gluon':
trainer.set_learning_rate(args.lr)
else:
opt.lr = args.lr
print('lr change to', args.lr)
break
#_cb(param)
if mbatch%1000==0:
print('lr-batch-epoch:',args.lr, mbatch)
if mbatch>0 and mbatch%args.verbose==0:
save_step[0]+=1
msave = save_step[0]
do_save = False
is_highest = False
if args.task=='age' or args.task=='gender':
acc = val_test(mbatch)
if acc>=highest_acc[-1]:
highest_acc[-1] = acc
is_highest = True
do_save = True
else:
acc_list = ver_test(mbatch)
if len(acc_list)>0:
lfw_score = acc_list[0]
if lfw_score>highest_acc[0]:
highest_acc[0] = lfw_score
if lfw_score>=0.998:
do_save = True
if acc_list[-1]>=highest_acc[-1]:
highest_acc[-1] = acc_list[-1]
if lfw_score>=0.99:
do_save = True
is_highest = True
if args.ckpt==0:
do_save = False
elif args.ckpt>1:
do_save = True
if do_save:
print('saving', msave)
#print('saving gluon params')
fname = os.path.join(args.prefix, 'model-gluon.params')
net.save_params(fname)
fname = os.path.join(args.prefix, 'model')
net.export(fname, msave)
#arg, aux = model.get_params()
#mx.model.save_checkpoint(prefix, msave, model.symbol, arg, aux)
print('[%d]Accuracy-Highest: %1.5f'%(mbatch, highest_acc[-1]))
if args.max_steps>0 and mbatch>args.max_steps:
sys.exit(0)
def _batch_callback_sym(param):
_cb(param)
_batch_callback()
if args.mode!='gluon':
model = mx.mod.Module(
context = ctx,
symbol = sym,
)
model.fit(train_iter,
begin_epoch = 0,
num_epoch = args.end_epoch,
eval_data = None,
eval_metric = metric,
kvstore = 'device',
optimizer = opt,
initializer = initializer,
arg_params = arg_params,
aux_params = aux_params,
allow_missing = True,
batch_end_callback = _batch_callback_sym,
epoch_end_callback = None )
else:
loss_weight = 1.0
if args.task=='age':
loss_weight = 1.0/AGE
#loss = gluon.loss.SoftmaxCrossEntropyLoss(weight = loss_weight)
loss = nd.SoftmaxOutput
#loss = gluon.loss.SoftmaxCrossEntropyLoss()
while True:
#trainer = update_learning_rate(opt.lr, trainer, epoch, opt.lr_factor, lr_steps)
tic = time.time()
train_iter.reset()
metric.reset()
btic = time.time()
for i, batch in enumerate(train_iter):
_batch_callback()
#data = gluon.utils.split_and_load(batch.data[0].astype(opt.dtype), ctx_list=ctx, batch_axis=0)
#label = gluon.utils.split_and_load(batch.label[0].astype(opt.dtype), ctx_list=ctx, batch_axis=0)
data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
outputs = []
Ls = []
with ag.record():
for x, y in zip(data, label):
#print(y.asnumpy())
if args.task=='':
if args.margin_a>0.0:
z = net(x,y)
else:
z = net(x)
#print(z[0].shape, z[1].shape)
else:
z = net(x)
if args.task=='gender':
L = loss(z[1], y)
#L = L/args.per_batch_size
Ls.append(L)
outputs.append(z[1])
elif args.task=='age':
for k in xrange(AGE):
_z = nd.slice_axis(z[2], axis=1, begin=k*2, end=k*2+2)
_y = nd.slice_axis(y, axis=1, begin=k, end=k+1)
_y = nd.flatten(_y)
L = loss(_z, _y)
#L = L/args.per_batch_size
#L /= AGE
Ls.append(L)
outputs.append(z[2])
else:
L = loss(z, y)
#L = L/args.per_batch_size
Ls.append(L)
outputs.append(z)
# store the loss and do backward after we have done forward
# on all GPUs for better speed on multiple GPUs.
ag.backward(Ls)
#trainer.step(batch.data[0].shape[0], ignore_stale_grad=True)
#trainer.step(args.ctx_num)
n = batch.data[0].shape[0]
#print(n,n)
trainer.step(n)
metric.update(label, outputs)
if i>0 and i%20==0:
name, acc = metric.get()
if len(name)==2:
logger.info('Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f, %s=%f'%(
num_epochs, i, args.batch_size/(time.time()-btic), name[0], acc[0], name[1], acc[1]))
else:
logger.info('Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f'%(
num_epochs, i, args.batch_size/(time.time()-btic), name[0], acc[0]))
#metric.reset()
btic = time.time()
epoch_time = time.time()-tic
# First epoch will usually be much slower than the subsequent epics,
# so don't factor into the average
if num_epochs > 0:
total_time = total_time + epoch_time
#name, acc = metric.get()
#logger.info('[Epoch %d] training: %s=%f, %s=%f'%(num_epochs, name[0], acc[0], name[1], acc[1]))
logger.info('[Epoch %d] time cost: %f'%(num_epochs, epoch_time))
num_epochs = num_epochs + 1
#name, val_acc = test(ctx, val_data)
#logger.info('[Epoch %d] validation: %s=%f, %s=%f'%(epoch, name[0], val_acc[0], name[1], val_acc[1]))
# save model if meet requirements
#save_checkpoint(epoch, val_acc[0], best_acc)
if num_epochs > 1:
print('Average epoch time: {}'.format(float(total_time)/(num_epochs - 1)))