def train_net(args):
ctx = []
cvd = os.environ['CUDA_VISIBLE_DEVICES'] ='0'
if len(cvd) > 0:
for i in range(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 = os.path.join(args.models_root,
'%s-%s-%s' % (args.network, args.loss, args.dataset),
'model')
prefix_dir = os.path.dirname(prefix)
print('prefix', prefix)
if not os.path.exists(prefix_dir):
os.makedirs(prefix_dir)
args.ctx_num = len(ctx)
args.batch_size = args.per_batch_size * args.ctx_num
args.rescale_threshold = 0
args.image_channel = config.image_shape[2]
config.batch_size = args.batch_size
config.per_batch_size = args.per_batch_size
data_dir = config.dataset_path
path_imgrec = None
path_imglist = None
image_size = config.image_shape[0:2]
assert len(image_size) == 2
assert image_size[0] == image_size[1]
print('image_size', image_size)
print('num_classes', config.num_classes)
path_imgrec = os.path.join(data_dir, "train.rec")
print('Called with argument:', args, config)
data_shape = (args.image_channel, image_size[0], image_size[1])
mean = None
begin_epoch = 0
if len(args.pretrained) == 0:
arg_params = None
aux_params = None
sym = get_symbol(args)
if config.net_name == 'spherenet':
data_shape_dict = {'data': (args.per_batch_size, ) + data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
else:
print('loading', args.pretrained, args.pretrained_epoch)
_, arg_params, aux_params = mx.model.load_checkpoint(
args.pretrained, args.pretrained_epoch)
sym = get_symbol(args)
if config.count_flops:
all_layers = sym.get_internals()
_sym = all_layers['fc1_output']
FLOPs = flops_counter.count_flops(_sym,
data=(1, 3, image_size[0],
image_size[1]))
_str = flops_counter.flops_str(FLOPs)
print('Network FLOPs: %s' % _str)
#label_name = 'softmax_label'
#label_shape = (args.batch_size,)
model = mx.mod.Module(
context=ctx,
symbol=sym,
)
val_dataiter = None
if config.loss_name.find('triplet') >= 0:
from triplet_image_iter import FaceImageIter
triplet_params = [
config.triplet_bag_size, config.triplet_alpha,
config.triplet_max_ap
]
train_dataiter = FaceImageIter(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=path_imgrec,
shuffle=True,
rand_mirror=config.data_rand_mirror,
mean=mean,
cutoff=config.data_cutoff,
ctx_num=args.ctx_num,
images_per_identity=config.images_per_identity,
triplet_params=triplet_params,
mx_model=model,
)
_metric = LossValueMetric()
eval_metrics = [mx.metric.create(_metric)]
else:
#from image_iter import FaceImageIter
#train_dataiter = FaceImageIter(
# batch_size=args.batch_size,
# data_shape=data_shape,
# path_imgrec=path_imgrec,
# shuffle=True,
# rand_mirror=config.data_rand_mirror,
# mean=mean,
# cutoff=config.data_cutoff,
# color_jittering=config.data_color,
# images_filter=config.data_images_filter,
#)
from image_iter import get_face_image_iter
train_dataiter = get_face_image_iter(config, data_shape, path_imgrec)
metric1 = AccMetric()
eval_metrics = [mx.metric.create(metric1)]
if config.ce_loss:
metric2 = LossValueMetric()
eval_metrics.append(mx.metric.create(metric2))
if config.net_name == 'fresnet' or config.net_name == 'fmobilefacenet':
initializer = mx.init.Xavier(rnd_type='gaussian',
factor_type="out",
magnitude=2) #resnet style
else:
initializer = mx.init.Xavier(rnd_type='uniform',
factor_type="in",
magnitude=2)
#initializer = mx.init.Xavier(rnd_type='gaussian', factor_type="out", magnitude=2) #resnet style
_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, args.frequent)
ver_list = []
ver_name_list = []
for name in config.val_targets:
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 range(len(ver_list)):
acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(
ver_list[i], model, args.batch_size, 10, None, None)
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
highest_acc = [0.0, 0.0] #lfw and target
#for i in range(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
lr_steps = [int(x) for x in args.lr_steps.split(',')]
print('lr_steps', lr_steps)
def _batch_callback(param):
#global global_step
global_step[0] += 1
mbatch = global_step[0]
for step in lr_steps:
if mbatch == step:
opt.lr *= 0.1
print('lr change to', opt.lr)
break
_cb(param)
if mbatch % 1000 == 0:
print('lr-batch-epoch:', opt.lr, param.nbatch, param.epoch)
if mbatch >= 0 and mbatch % args.verbose == 0:
acc_list = ver_test(mbatch)
save_step[0] += 1
msave = save_step[0]
do_save = False
is_highest = False
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
score = sum(acc_list)
if acc_list[-1] >= highest_acc[-1]:
if acc_list[-1] > highest_acc[-1]:
is_highest = True
else:
if score >= highest_acc[0]:
is_highest = True
highest_acc[0] = score
highest_acc[-1] = acc_list[-1]
#if lfw_score>=0.99:
# do_save = True
if is_highest:
do_save = True
if args.ckpt == 0:
do_save = False
elif args.ckpt == 2:
do_save = True
elif args.ckpt == 3:
msave = 1
if do_save:
print('saving', msave)
arg, aux = model.get_params()
if config.ckpt_embedding:
all_layers = model.symbol.get_internals()
_sym = all_layers['fc1_output']
_arg = {}
for k in arg:
if not k.startswith('fc7'):
_arg[k] = arg[k]
mx.model.save_checkpoint(prefix, msave, _sym, _arg, aux)
else:
mx.model.save_checkpoint(prefix, msave, model.symbol, arg,
aux)
print('[%d]Accuracy-Highest: %1.5f' % (mbatch, highest_acc[-1]))
if config.max_steps > 0 and mbatch > config.max_steps:
sys.exit(0)
epoch_cb = None
model.fit(
train_dataiter,
begin_epoch=begin_epoch,
num_epoch=999999,
eval_data=val_dataiter,
eval_metric=eval_metrics,
kvstore=args.kvstore,
optimizer=opt,
#optimizer_params = optimizer_params,
initializer=initializer,
arg_params=arg_params,
aux_params=aux_params,
allow_missing=True,
batch_end_callback=_batch_callback,
epoch_end_callback=epoch_cb)
def train_net(args):
#_seed = 727
#random.seed(_seed)
#np.random.seed(_seed)
#mx.random.seed(_seed)
config.fp16 = False
config.fp16_scale = 0.0
if args.fp16_scale > 0.0:
config.fp16 = True
config.fp16_scale = args.fp16_scale
print('use fp16, scale=', config.fp16_scale)
ctx = []
cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
if len(cvd) > 0:
for i in range(len(cvd.split(','))):
ctx.append(mx.gpu(i))
if len(ctx) == 0:
ctx = [mx.cpu()]
print('use cpu')
else:
print('gpu num:', len(ctx))
if len(args.extra_model_name) == 0:
prefix = os.path.join(
args.models_root,
'%s-%s-%s' % (args.network, args.loss, args.dataset), 'model')
else:
prefix = os.path.join(
args.models_root, '%s-%s-%s-%s' %
(args.network, args.loss, args.dataset, args.extra_model_name),
'model')
prefix_dir = os.path.dirname(prefix)
print('prefix', prefix)
if not os.path.exists(prefix_dir):
os.makedirs(prefix_dir)
args.ctx_num = len(ctx)
if args.per_batch_size == 0:
args.per_batch_size = 128
args.batch_size = args.per_batch_size * args.ctx_num
args.rescale_threshold = 0
args.image_channel = config.image_shape[2]
config.batch_size = args.batch_size
config.per_batch_size = args.per_batch_size
data_dir = config.dataset_path
path_imgrec = None
path_imglist = None
image_size = config.image_shape[0:2]
assert len(image_size) == 2
assert image_size[0] == image_size[1]
print('image_size', image_size)
print('num_classes', config.num_classes)
path_imgrec = os.path.join(data_dir, "train.rec")
data_shape = (args.image_channel, image_size[0], image_size[1])
num_workers = config.num_workers
global_num_ctx = num_workers * args.ctx_num
if config.num_classes % global_num_ctx == 0:
args.ctx_num_classes = config.num_classes // global_num_ctx
else:
args.ctx_num_classes = config.num_classes // global_num_ctx + 1
args.local_num_classes = args.ctx_num_classes * args.ctx_num
args.local_class_start = args.local_num_classes * args.worker_id
#if len(args.partial)==0:
# local_classes_range = (0, args.num_classes)
#else:
# _vec = args.partial.split(',')
# local_classes_range = (int(_vec[0]), int(_vec[1]))
#args.partial_num_classes = local_classes_range[1] - local_classes_range[0]
#args.partial_start = local_classes_range[0]
print('Called with argument:', args, config)
mean = None
begin_epoch = 0
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
arg_params = None
aux_params = None
if len(args.pretrained) == 0:
esym = get_symbol_embedding()
asym = get_symbol_arcface
else:
#assert False
print('loading', args.pretrained, args.pretrained_epoch)
pretrain_esym, arg_params, aux_params = mx.model.load_checkpoint(
args.pretrained, args.pretrained_epoch)
esym = get_symbol_embedding(pretrain_esym)
asym = get_symbol_arcface
if config.count_flops:
all_layers = esym.get_internals()
_sym = all_layers['fc1_output']
FLOPs = flops_counter.count_flops(_sym,
data=(1, 3, image_size[0],
image_size[1]))
_str = flops_counter.flops_str(FLOPs)
print('Network FLOPs: %s' % _str)
if config.num_workers == 1:
from parall_module_local_v1 import ParallModule
else:
from parall_module_dist import ParallModule
model = ParallModule(
context=ctx,
symbol=esym,
data_names=['data'],
label_names=['softmax_label'],
asymbol=asym,
args=args,
)
val_dataiter = None
if config.net_name == 'fresnet' or config.net_name == 'fmobilefacenet':
initializer = mx.init.Xavier(rnd_type='gaussian',
factor_type="out",
magnitude=2) #resnet style
else:
initializer = mx.init.Xavier(rnd_type='uniform',
factor_type="in",
magnitude=2)
_rescale = 1.0 / args.batch_size
if config.fp16:
opt = optimizer.SGD(learning_rate=base_lr,
momentum=base_mom,
wd=base_wd,
rescale_grad=_rescale,
multi_precision=True)
else:
opt = optimizer.SGD(learning_rate=base_lr,
momentum=base_mom,
wd=base_wd,
rescale_grad=_rescale,
multi_precision=False)
opt_fc7 = optimizer.SGD(learning_rate=base_lr,
momentum=base_mom,
wd=base_wd,
rescale_grad=_rescale,
multi_precision=False)
_cb = mx.callback.Speedometer(args.batch_size, args.frequent)
ver_list = []
ver_name_list = []
for name in config.val_targets:
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 range(len(ver_list)):
acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(
ver_list[i], model, args.batch_size, 10, None, None)
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
highest_acc = [0.0, 0.0] #lfw and target
#for i in range(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
lr_steps = [int(x) for x in args.lr_steps.split(',')]
print('lr_steps', lr_steps)
def _batch_callback(param):
#global global_step
global_step[0] += 1
mbatch = global_step[0]
for step in lr_steps:
if mbatch == step:
opt.lr *= 0.1
opt_fc7.lr *= 0.1
print('lr change to', opt.lr, opt_fc7.lr)
break
_cb(param)
if mbatch % 1000 == 0:
print('lr-batch-epoch:', opt.lr, opt_fc7.lr, param.nbatch,
param.epoch)
if mbatch >= 0 and mbatch % args.verbose == 0:
acc_list = ver_test(mbatch)
save_step[0] += 1
msave = save_step[0]
do_save = False
is_highest = False
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
score = sum(acc_list)
if acc_list[-1] >= highest_acc[-1]:
if acc_list[-1] > highest_acc[-1]:
is_highest = True
else:
if score >= highest_acc[0]:
is_highest = True
highest_acc[0] = score
highest_acc[-1] = acc_list[-1]
#if lfw_score>=0.99:
# do_save = True
if is_highest:
do_save = True
if args.ckpt == 0:
do_save = False
elif args.ckpt == 2:
do_save = True
elif args.ckpt == 3:
msave = 1
if do_save:
print('saving', msave)
arg, aux = model.get_export_params()
all_layers = model.symbol.get_internals()
_sym = all_layers['fc1_output']
mx.model.save_checkpoint(prefix, msave, _sym, arg, aux)
print('[%d]Accuracy-Highest: %1.5f' % (mbatch, highest_acc[-1]))
if config.max_steps > 0 and mbatch > config.max_steps:
msave = 0
config.fp16 = False
print('saving last', msave)
arg, aux = model.get_export_params()
_sym = eval(config.net_name).get_symbol()
mx.model.save_checkpoint(prefix, msave, _sym, arg, aux)
sys.exit(0)
epoch_cb = None
train_dataiter = get_face_image_iter(config, data_shape, path_imgrec)
#train_dataiter = FaceImageIter(
# batch_size=args.batch_size,
# data_shape=data_shape,
# path_imgrec=path_imgrec,
# shuffle=True,
# rand_mirror=config.data_rand_mirror,
# mean=mean,
# cutoff=config.data_cutoff,
# color_jittering=config.data_color,
# images_filter=config.data_images_filter,
#)
#train_dataiter = mx.io.PrefetchingIter(train_dataiter)
model.fit(
train_dataiter,
begin_epoch=begin_epoch,
num_epoch=999999,
eval_data=val_dataiter,
#eval_metric = eval_metrics,
kvstore=args.kvstore,
optimizer=[opt, opt_fc7],
#optimizer_params = optimizer_params,
initializer=initializer,
arg_params=arg_params,
aux_params=aux_params,
allow_missing=True,
batch_end_callback=_batch_callback,
epoch_end_callback=epoch_cb)