def main(args):
ctx = mx.gpu(args.gpu)
args.ctx_num = 1
prop = face_image.load_property(args.data)
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading',prefix, epoch)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
arg_params, aux_params = ch_dev(arg_params, aux_params, ctx)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
#model = mx.mod.Module.load(prefix, epoch, context = ctx)
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
#model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))])
model.set_params(arg_params, aux_params)
path_imgrec = os.path.join(args.data, 'train.rec')
path_imgidx = os.path.join(args.data, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag>0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
stat = []
count = 0
data = nd.zeros( (1 ,3, image_size[0], image_size[1]) )
label = nd.zeros( (1,) )
for idx in imgidx:
if len(stat)%100==0:
print('processing', len(stat))
s = imgrec.read_idx(idx)
header, img = mx.recordio.unpack(s)
img = mx.image.imdecode(img)
img = nd.transpose(img, axes=(2, 0, 1))
data[0][:] = img
#input_blob = np.expand_dims(img.asnumpy(), axis=0)
#arg_params["data"] = mx.nd.array(input_blob, ctx)
#arg_params["softmax_label"] = mx.nd.empty((1,), ctx)
time_now = datetime.datetime.now()
#exe = sym.bind(ctx, arg_params ,args_grad=None, grad_req="null", aux_states=aux_params)
#exe.forward(is_train=False)
#_embedding = exe.outputs[0].asnumpy().flatten()
#db = mx.io.DataBatch(data=(data,), label=(label,))
db = mx.io.DataBatch(data=(data,))
model.forward(db, is_train=False)
net_out = model.get_outputs()[0].asnumpy()
time_now2 = datetime.datetime.now()
diff = time_now2 - time_now
stat.append(diff.total_seconds())
if len(stat)==args.param1:
break
stat = stat[10:]
print('avg infer time', np.mean(stat))
def main(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))
ctx_num = len(ctx)
path_imgrec = os.path.join(args.input, 'train.rec')
path_imgidx = os.path.join(args.input, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
outf = open(os.path.join(args.input, 'c2c'), 'w')
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag>0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
id2range = {}
seq_identity = range(int(header.label[0]), int(header.label[1]))
for identity in seq_identity:
s = imgrec.read_idx(identity)
header, _ = mx.recordio.unpack(s)
id2range[identity] = (int(header.label[0]), int(header.label[1]))
print('id2range', len(id2range))
prop = face_image.load_property(args.input)
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading',prefix, epoch)
model = mx.mod.Module.load(prefix, epoch, context = ctx)
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
nrof_images = 0
nrof_removed = 0
idx = 1
id2label = {}
pp = 0
for _id, v in id2range.iteritems():
pp+=1
if pp%100==0:
print('processing id', pp)
_list = range(*v)
ocontents = []
for i in xrange(len(_list)):
_idx = _list[i]
#print('_idx', _id, _idx)
s = imgrec.read_idx(_idx)
ocontents.append(s)
#continue
embeddings = None
headers = [None]*len(ocontents)
#print(len(ocontents))
ba = 0
while True:
bb = min(ba+args.batch_size, len(ocontents))
if ba>=bb:
break
_batch_size = bb-ba
_batch_size2 = max(_batch_size, ctx_num)
data = nd.zeros( (_batch_size2,3, image_size[0], image_size[1]) )
label = nd.zeros( (_batch_size2,) )
count = bb-ba
ii=0
for i in xrange(ba, bb):
header, img = mx.recordio.unpack(ocontents[i])
headers[i] = header
img = mx.image.imdecode(img)
img = nd.transpose(img, axes=(2, 0, 1))
data[ii][:] = img
label[ii][:] = header.label[0]
ii+=1
while ii<_batch_size2:
data[ii][:] = data[0][:]
label[ii][:] = label[0][:]
ii+=1
db = mx.io.DataBatch(data=(data,), label=(label,))
model.forward(db, is_train=False)
net_out = model.get_outputs()
net_out = net_out[0].asnumpy()
if embeddings is None:
embeddings = np.zeros( (len(ocontents), net_out.shape[1]))
embeddings[ba:bb,:] = net_out[0:_batch_size,:]
ba = bb
embeddings = sklearn.preprocessing.normalize(embeddings)
emb_mean = np.mean(embeddings, axis=0, keepdims=True)
emb_mean = sklearn.preprocessing.normalize(emb_mean)
sim = np.dot(embeddings, emb_mean.T)
#print(sim.shape)
sims = sim.flatten()
assert len(_list)==len(sims)
assert len(_list)==len(ocontents)
for i in xrange(len(ocontents)):
_sim = sims[i]
_idx = _list[i]
_header = headers[i]
#TODO
outf.write("%d,%f,%d\n"%(_idx, _sim, int(_header.label[1])))
outf.close()
def do_clean(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))
ctx_num = len(ctx)
path_imgrec = os.path.join(args.input, 'train.rec')
path_imgidx = os.path.join(args.input, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag>0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
id2range = {}
seq_identity = range(int(header.label[0]), int(header.label[1]))
for identity in seq_identity:
s = imgrec.read_idx(identity)
header, _ = mx.recordio.unpack(s)
id2range[identity] = (int(header.label[0]), int(header.label[1]))
print('id2range', len(id2range))
prop = face_image.load_property(args.input)
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading',prefix, epoch)
model = mx.mod.Module.load(prefix, epoch, context = ctx)
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
if args.test==0:
if not os.path.exists(args.output):
os.makedirs(args.output)
writer = mx.recordio.MXIndexedRecordIO(os.path.join(args.output, 'train.idx'), os.path.join(args.output, 'train.rec'), 'w')
nrof_images = 0
nrof_removed = 0
idx = 1
id2label = {}
pp = 0
for _id, v in id2range.iteritems():
pp+=1
if pp%100==0:
print('stat', nrof_images, nrof_removed)
_list = range(*v)
ocontents = []
for i in xrange(len(_list)):
_idx = _list[i]
s = imgrec.read_idx(_idx)
ocontents.append(s)
if len(ocontents)>15:
nrof_removed+=len(ocontents)
continue
embeddings = None
#print(len(ocontents))
ba = 0
while True:
bb = min(ba+args.batch_size, len(ocontents))
if ba>=bb:
break
_batch_size = bb-ba
_batch_size2 = max(_batch_size, ctx_num)
data = nd.zeros( (_batch_size2,3, image_size[0], image_size[1]) )
label = nd.zeros( (_batch_size2,) )
count = bb-ba
ii=0
for i in xrange(ba, bb):
header, img = mx.recordio.unpack(ocontents[i])
img = mx.image.imdecode(img)
img = nd.transpose(img, axes=(2, 0, 1))
data[ii][:] = img
label[ii][:] = header.label
ii+=1
while ii<_batch_size2:
data[ii][:] = data[0][:]
label[ii][:] = label[0][:]
ii+=1
db = mx.io.DataBatch(data=(data,), label=(label,))
model.forward(db, is_train=False)
net_out = model.get_outputs()
net_out = net_out[0].asnumpy()
if embeddings is None:
embeddings = np.zeros( (len(ocontents), net_out.shape[1]))
embeddings[ba:bb,:] = net_out[0:_batch_size,:]
ba = bb
embeddings = sklearn.preprocessing.normalize(embeddings)
contents = []
if args.mode==1:
emb_mean = np.mean(embeddings, axis=0, keepdims=True)
emb_mean = sklearn.preprocessing.normalize(emb_mean)
sim = np.dot(embeddings, emb_mean.T)
#print(sim.shape)
sim = sim.flatten()
#print(sim.flatten())
x = np.argsort(sim)
for ix in xrange(len(x)):
_idx = x[ix]
_sim = sim[_idx]
#if ix<int(len(x)*0.3) and _sim<args.threshold:
if _sim<args.threshold:
continue
contents.append(ocontents[_idx])
else:
y_pred = DBSCAN(eps = args.threshold, min_samples = 2).fit_predict(embeddings)
#print(y_pred)
gmap = {}
for _idx in xrange(embeddings.shape[0]):
label = int(y_pred[_idx])
if label not in gmap:
gmap[label] = []
gmap[label].append(_idx)
assert len(gmap)>0
_max = [0, 0]
for label in xrange(10):
if not label in gmap:
break
glist = gmap[label]
if len(glist)>_max[1]:
_max[0] = label
_max[1] = len(glist)
if _max[1]>0:
glist = gmap[_max[0]]
for _idx in glist:
contents.append(ocontents[_idx])
nrof_removed+=(len(ocontents)-len(contents))
if len(contents)==0:
continue
#assert len(contents)>0
id2label[_id] = (idx, idx+len(contents))
nrof_images += len(contents)
for content in contents:
if args.test==0:
writer.write_idx(idx, content)
idx+=1
id_idx = idx
if args.test==0:
for _id, _label in id2label.iteritems():
_header = mx.recordio.IRHeader(1, _label, idx, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(idx, s)
idx+=1
_header = mx.recordio.IRHeader(1, (id_idx, idx), 0, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(0, s)
print(nrof_images, nrof_removed)
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)))
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
if args.loss_type==10:
args.per_batch_size = 256
args.batch_size = args.per_batch_size*args.ctx_num
args.rescale_threshold = 0
args.image_channel = 3
ppatch = [int(x) for x in args.patch.split('_')]
assert len(ppatch)==5
os.environ['BETA'] = str(args.beta)
data_dir_list = args.data_dir.split(',')
if args.loss_type!=12 and args.loss_type!=13:
assert len(data_dir_list)==1
data_dir = data_dir_list[0]
args.use_val = False
path_imgrec = None
path_imglist = None
val_rec = 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)
args.coco_scale = 0.5*math.log(float(args.num_classes-1))+3
#path_imglist = "/raid5data/dplearn/MS-Celeb-Aligned/lst2"
path_imgrec = os.path.join(data_dir, "train.rec")
val_rec = os.path.join(data_dir, "val.rec")
if os.path.exists(val_rec) and args.loss_type<10:
args.use_val = True
else:
val_rec = None
#args.use_val = False
if args.loss_type==1 and args.num_classes>20000:
args.beta_freeze = 5000
args.gamma = 0.06
if args.loss_type<9:
assert args.images_per_identity==0
else:
if args.images_per_identity==0:
if args.loss_type==11:
args.images_per_identity = 2
elif args.loss_type==10 or args.loss_type==9:
args.images_per_identity = 16
elif args.loss_type==12 or args.loss_type==13:
args.images_per_identity = 5
assert args.per_batch_size%3==0
assert args.images_per_identity>=2
args.per_identities = int(args.per_batch_size/args.images_per_identity)
print('Called with argument:', args)
data_shape = (args.image_channel,image_size[0],image_size[1])
mean = None
begin_epoch = 0
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained)==0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(',')
print('loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
if args.network[0]=='s':
data_shape_dict = {'data' : (args.per_batch_size,)+data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
data_extra = None
hard_mining = False
triplet_params = None
coco_mode = False
if args.loss_type==10:
hard_mining = True
_shape = (args.batch_size, args.per_batch_size)
data_extra = np.full(_shape, -1.0, dtype=np.float32)
c = 0
while c<args.batch_size:
a = 0
while a<args.per_batch_size:
b = a+args.images_per_identity
data_extra[(c+a):(c+b),a:b] = 1.0
#print(c+a, c+b, a, b)
a = b
c += args.per_batch_size
elif args.loss_type==11:
data_extra = np.zeros( (args.batch_size, args.per_identities), dtype=np.float32)
c = 0
while c<args.batch_size:
for i in xrange(args.per_identities):
data_extra[c+i][i] = 1.0
c+=args.per_batch_size
elif args.loss_type==12 or args.loss_type==13:
triplet_params = [args.triplet_bag_size, args.triplet_alpha, args.triplet_max_ap]
elif args.loss_type==9:
coco_mode = True
label_name = 'softmax_label'
label_shape = (args.batch_size,)
if args.output_c2c:
label_shape = (args.batch_size,2)
if data_extra is None:
model = mx.mod.Module(
context = ctx,
symbol = sym,
)
else:
data_names = ('data', 'extra')
#label_name = ''
model = mx.mod.Module(
context = ctx,
symbol = sym,
data_names = data_names,
label_names = (label_name,),
)
if args.use_val:
val_dataiter = FaceImageIter(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = val_rec,
#path_imglist = val_path,
shuffle = False,
rand_mirror = False,
mean = mean,
ctx_num = args.ctx_num,
data_extra = data_extra,
)
else:
val_dataiter = None
if len(data_dir_list)==1 and args.loss_type!=12 and args.loss_type!=13:
train_dataiter = 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,
c2c_threshold = args.c2c_threshold,
output_c2c = args.output_c2c,
c2c_mode = args.c2c_mode,
limit = args.train_limit,
ctx_num = args.ctx_num,
images_per_identity = args.images_per_identity,
data_extra = data_extra,
hard_mining = hard_mining,
triplet_params = triplet_params,
coco_mode = coco_mode,
mx_model = model,
label_name = label_name,
)
else:
iter_list = []
for _data_dir in data_dir_list:
_path_imgrec = os.path.join(_data_dir, "train.rec")
_dataiter = 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,
c2c_threshold = args.c2c_threshold,
output_c2c = args.output_c2c,
c2c_mode = args.c2c_mode,
limit = args.train_limit,
ctx_num = args.ctx_num,
images_per_identity = args.images_per_identity,
data_extra = data_extra,
hard_mining = hard_mining,
triplet_params = triplet_params,
coco_mode = coco_mode,
mx_model = model,
label_name = label_name,
)
iter_list.append(_dataiter)
iter_list.append(_dataiter)
train_dataiter = FaceImageIterList(iter_list)
if args.loss_type<10:
_metric = AccMetric()
else:
_metric = LossValueMetric()
eval_metrics = [mx.metric.create(_metric)]
if args.network[0]=='r':
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)
_rescale = 1.0/args.ctx_num
if args.noise_sgd>0.0:
print('use noise sgd')
opt = NoiseSGD(scale = args.noise_sgd, learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
else:
opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
som = 20
if args.loss_type==12 or args.loss_type==13:
som = 2
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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], model, args.batch_size, 10, data_extra, label_shape)
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():
acc = AccMetric()
val_metric = mx.metric.create(acc)
val_metric.reset()
val_dataiter.reset()
for i, eval_batch in enumerate(val_dataiter):
model.forward(eval_batch, is_train=False)
model.update_metric(val_metric, eval_batch.label)
acc_value = val_metric.get_name_value()[0][1]
print('VACC: %f'%(acc_value))
highest_acc = [0.0, 0.0] #lfw and target
#for i in xrange(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
if len(args.lr_steps)==0:
lr_steps = [40000, 60000, 80000]
if args.loss_type>=1 and args.loss_type<=7:
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)
def _batch_callback(param):
#global global_step
global_step[0]+=1
mbatch = global_step[0]
for _lr in lr_steps:
if mbatch==args.beta_freeze+_lr:
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
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
if args.ckpt==0:
do_save = False
elif args.ckpt>1:
do_save = True
#for i in xrange(len(acc_list)):
# acc = acc_list[i]
# if acc>=highest_acc[i]:
# highest_acc[i] = acc
# if lfw_score>=0.99:
# do_save = True
#if args.loss_type==1 and mbatch>lr_steps[-1] and mbatch%10000==0:
# do_save = True
if do_save:
print('saving', msave)
if val_dataiter is not None:
val_test()
arg, aux = model.get_params()
mx.model.save_checkpoint(prefix, msave, model.symbol, arg, aux)
#if acc>=highest_acc[0]:
# lfw_npy = "%s-lfw-%04d" % (prefix, msave)
# X = np.concatenate(embeddings_list, axis=0)
# print('saving lfw npy', X.shape)
# np.save(lfw_npy, X)
print('[%d]Accuracy-Highest: %1.5f'%(mbatch, highest_acc[-1]))
if mbatch<=args.beta_freeze:
_beta = args.beta
else:
move = max(0, mbatch-args.beta_freeze)
_beta = max(args.beta_min, args.beta*math.pow(1+args.gamma*move, -1.0*args.power))
#print('beta', _beta)
os.environ['BETA'] = str(_beta)
if args.max_steps>0 and mbatch>args.max_steps:
sys.exit(0)
#epoch_cb = mx.callback.do_checkpoint(prefix, 1)
epoch_cb = None
#def _epoch_callback(epoch, sym, arg_params, aux_params):
# print('epoch-end', epoch)
model.fit(train_dataiter,
begin_epoch = begin_epoch,
num_epoch = end_epoch,
eval_data = val_dataiter,
eval_metric = eval_metrics,
kvstore = 'device',
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 )
args = parser.parse_args()
args.prefix = os.path.abspath(args.prefix)
#args.root = os.path.abspath(args.root)
return args
if __name__ == '__main__':
args = parse_args()
if args.list:
pass
#make_list(args)
else:
if os.path.isdir(args.prefix):
working_dir = args.prefix
else:
working_dir = os.path.dirname(args.prefix)
prop = face_image.load_property(working_dir)
image_size = prop.image_size
print('image_size', image_size)
args.image_h = image_size[0]
args.image_w = image_size[1]
files = [os.path.join(working_dir, fname) for fname in os.listdir(working_dir)
if os.path.isfile(os.path.join(working_dir, fname))]
count = 0
for fname in files:
if fname.startswith(args.prefix) and fname.endswith('.lst'):
print('Creating .rec file from', fname, 'in', working_dir)
count += 1
image_list = read_list(fname)
# -- write_record -- #
if args.num_thread > 1 and multiprocessing is not None:
q_in = [multiprocessing.Queue(1024) for i in range(args.num_thread)]
def main(args):
# ctx = mx.gpu(args.gpu)
ctx = mx.gpu(args.gpu) if mx.context.num_gpus() else mx.cpu(args.gpu)
args.ctx_num = 1
prop = face_image.load_property(args.data)
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading', prefix, epoch)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
arg_params, aux_params = ch_dev(arg_params, aux_params, ctx)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
#model = mx.mod.Module.load(prefix, epoch, context = ctx)
model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
#model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0],
image_size[1]))])
model.set_params(arg_params, aux_params)
path_imgrec = os.path.join(args.data, 'train.rec')
path_imgidx = os.path.join(args.data, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag > 0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
stat = []
count = 0
data = nd.zeros((1, 3, image_size[0], image_size[1]))
label = nd.zeros((1, ))
for idx in imgidx:
if len(stat) % 100 == 0:
print('processing', len(stat))
s = imgrec.read_idx(idx)
header, img = mx.recordio.unpack(s)
img = mx.image.imdecode(img)
img = nd.transpose(img, axes=(2, 0, 1))
data[0][:] = img
#input_blob = np.expand_dims(img.asnumpy(), axis=0)
#arg_params["data"] = mx.nd.array(input_blob, ctx)
#arg_params["softmax_label"] = mx.nd.empty((1,), ctx)
time_now = datetime.datetime.now()
#exe = sym.bind(ctx, arg_params ,args_grad=None, grad_req="null", aux_states=aux_params)
#exe.forward(is_train=False)
#_embedding = exe.outputs[0].asnumpy().flatten()
#db = mx.io.DataBatch(data=(data,), label=(label,))
db = mx.io.DataBatch(data=(data, ))
model.forward(db, is_train=False)
net_out = model.get_outputs()[0].asnumpy()
time_now2 = datetime.datetime.now()
diff = time_now2 - time_now
stat.append(diff.total_seconds())
if len(stat) == args.param1:
break
stat = stat[10:]
print('avg infer time', np.mean(stat))
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
if args.loss_type==10:
args.per_batch_size = 256
args.batch_size = args.per_batch_size*args.ctx_num
args.rescale_threshold = 0
args.image_channel = 3
ppatch = [int(x) for x in args.patch.split('_')]
assert len(ppatch)==5
os.environ['BETA'] = str(args.beta)
data_dir_list = args.data_dir.split(',')
if args.loss_type!=12 and args.loss_type!=13:
assert len(data_dir_list)==1
data_dir = data_dir_list[0]
args.use_val = False
path_imgrec = None
path_imglist = None
val_rec = 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)
args.coco_scale = 0.5*math.log(float(args.num_classes-1))+3
#path_imglist = "/raid5data/dplearn/MS-Celeb-Aligned/lst2"
path_imgrec = os.path.join(data_dir, "train.rec")
val_rec = os.path.join(data_dir, "val.rec")
if os.path.exists(val_rec) and args.loss_type<10:
args.use_val = True
else:
val_rec = None
#args.use_val = False
if args.loss_type==1 and args.num_classes>20000:
args.beta_freeze = 5000
args.gamma = 0.06
if args.loss_type<9:
assert args.images_per_identity==0
else:
if args.images_per_identity==0:
if args.loss_type==11:
args.images_per_identity = 2
elif args.loss_type==10 or args.loss_type==9:
args.images_per_identity = 16
elif args.loss_type==12 or args.loss_type==13:
args.images_per_identity = 5
assert args.per_batch_size%3==0
assert args.images_per_identity>=2
args.per_identities = int(args.per_batch_size/args.images_per_identity)
print('Called with argument:', args)
data_shape = (args.image_channel,image_size[0],image_size[1])
mean = None
begin_epoch = 0
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained)==0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(',')
print('loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
if args.network[0]=='s':
data_shape_dict = {'data' : (args.per_batch_size,)+data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
data_extra = None
hard_mining = False
triplet_params = None
coco_mode = False
if args.loss_type==10:
hard_mining = True
_shape = (args.batch_size, args.per_batch_size)
data_extra = np.full(_shape, -1.0, dtype=np.float32)
c = 0
while c<args.batch_size:
a = 0
while a<args.per_batch_size:
b = a+args.images_per_identity
data_extra[(c+a):(c+b),a:b] = 1.0
#print(c+a, c+b, a, b)
a = b
c += args.per_batch_size
elif args.loss_type==11:
data_extra = np.zeros( (args.batch_size, args.per_identities), dtype=np.float32)
c = 0
while c<args.batch_size:
for i in xrange(args.per_identities):
data_extra[c+i][i] = 1.0
c+=args.per_batch_size
elif args.loss_type==12 or args.loss_type==13:
triplet_params = [args.triplet_bag_size, args.triplet_alpha, args.triplet_max_ap]
elif args.loss_type==9:
coco_mode = True
label_name = 'softmax_label'
label_shape = (args.batch_size,)
if args.output_c2c:
label_shape = (args.batch_size,2)
if data_extra is None:
model = mx.mod.Module(
context = ctx,
symbol = sym,
)
else:
data_names = ('data', 'extra')
#label_name = ''
model = mx.mod.Module(
context = ctx,
symbol = sym,
data_names = data_names,
label_names = (label_name,),
)
if args.use_val:
val_dataiter = FaceImageIter(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = val_rec,
#path_imglist = val_path,
shuffle = False,
rand_mirror = False,
mean = mean,
ctx_num = args.ctx_num,
data_extra = data_extra,
)
else:
val_dataiter = None
if len(data_dir_list)==1 and args.loss_type!=12 and args.loss_type!=13:
train_dataiter = 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,
c2c_threshold = args.c2c_threshold,
output_c2c = args.output_c2c,
c2c_mode = args.c2c_mode,
limit = args.train_limit,
ctx_num = args.ctx_num,
images_per_identity = args.images_per_identity,
data_extra = data_extra,
hard_mining = hard_mining,
triplet_params = triplet_params,
coco_mode = coco_mode,
mx_model = model,
label_name = label_name,
)
else:
iter_list = []
for _data_dir in data_dir_list:
_path_imgrec = os.path.join(_data_dir, "train.rec")
_dataiter = 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,
c2c_threshold = args.c2c_threshold,
output_c2c = args.output_c2c,
c2c_mode = args.c2c_mode,
limit = args.train_limit,
ctx_num = args.ctx_num,
images_per_identity = args.images_per_identity,
data_extra = data_extra,
hard_mining = hard_mining,
triplet_params = triplet_params,
coco_mode = coco_mode,
mx_model = model,
label_name = label_name,
)
iter_list.append(_dataiter)
iter_list.append(_dataiter)
train_dataiter = FaceImageIterList(iter_list)
if args.loss_type<10:
_metric = AccMetric()
else:
_metric = LossValueMetric()
eval_metrics = [mx.metric.create(_metric)]
if args.network[0]=='r':
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)
_rescale = 1.0/args.ctx_num
if args.noise_sgd>0.0:
print('use noise sgd')
opt = NoiseSGD(scale = args.noise_sgd, learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
else:
opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
som = 20
if args.loss_type==12 or args.loss_type==13:
som = 2
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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], model, args.batch_size, 10, data_extra, label_shape)
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():
acc = AccMetric()
val_metric = mx.metric.create(acc)
val_metric.reset()
val_dataiter.reset()
for i, eval_batch in enumerate(val_dataiter):
model.forward(eval_batch, is_train=False)
model.update_metric(val_metric, eval_batch.label)
acc_value = val_metric.get_name_value()[0][1]
print('VACC: %f'%(acc_value))
highest_acc = [0.0, 0.0] #lfw and target
#for i in xrange(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
if len(args.lr_steps)==0:
lr_steps = [40000, 60000, 80000]
if args.loss_type>=1 and args.loss_type<=7:
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)
def _batch_callback(param):
#global global_step
global_step[0]+=1
mbatch = global_step[0]
for _lr in lr_steps:
if mbatch==args.beta_freeze+_lr:
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
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
if args.ckpt==0:
do_save = False
elif args.ckpt>1:
do_save = True
#for i in xrange(len(acc_list)):
# acc = acc_list[i]
# if acc>=highest_acc[i]:
# highest_acc[i] = acc
# if lfw_score>=0.99:
# do_save = True
#if args.loss_type==1 and mbatch>lr_steps[-1] and mbatch%10000==0:
# do_save = True
if do_save:
print('saving', msave)
if val_dataiter is not None:
val_test()
arg, aux = model.get_params()
mx.model.save_checkpoint(prefix, msave, model.symbol, arg, aux)
#if acc>=highest_acc[0]:
# lfw_npy = "%s-lfw-%04d" % (prefix, msave)
# X = np.concatenate(embeddings_list, axis=0)
# print('saving lfw npy', X.shape)
# np.save(lfw_npy, X)
print('[%d]Accuracy-Highest: %1.5f'%(mbatch, highest_acc[-1]))
if mbatch<=args.beta_freeze:
_beta = args.beta
else:
move = max(0, mbatch-args.beta_freeze)
_beta = max(args.beta_min, args.beta*math.pow(1+args.gamma*move, -1.0*args.power))
#print('beta', _beta)
os.environ['BETA'] = str(_beta)
if args.max_steps>0 and mbatch>args.max_steps:
sys.exit(0)
#epoch_cb = mx.callback.do_checkpoint(prefix, 1)
epoch_cb = None
#def _epoch_callback(epoch, sym, arg_params, aux_params):
# print('epoch-end', epoch)
model.fit(train_dataiter,
begin_epoch = begin_epoch,
num_epoch = end_epoch,
eval_data = val_dataiter,
eval_metric = eval_metrics,
kvstore = 'device',
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 )
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='do verification')
# general
parser.add_argument('--data-dir', default='', help='')
parser.add_argument('--model', default='../model/softmax,50', help='path to load model.')
parser.add_argument('--target', default='lfw,cfp_ff,cfp_fp,agedb_30', help='test targets.')
parser.add_argument('--gpu', default=0, type=int, help='gpu id')
parser.add_argument('--batch-size', default=32, type=int, help='')
parser.add_argument('--max', default='', type=str, help='')
parser.add_argument('--mode', default=0, type=int, help='')
parser.add_argument('--nfolds', default=10, type=int, help='')
args = parser.parse_args()
prop = face_image.load_property(args.data_dir)
image_size = prop.image_size
print('image_size', image_size)
ctx = mx.gpu(args.gpu)
nets = []
vec = args.model.split(',')
prefix = args.model.split(',')[0]
epochs = []
if len(vec)==1:
pdir = os.path.dirname(prefix)
for fname in os.listdir(pdir):
if not fname.endswith('.params'):
continue
_file = os.path.join(pdir, fname)
if _file.startswith(prefix):
epoch = int(fname.split('.')[0].split('-')[1])
#args.root = os.path.abspath(args.root)
return args
if __name__ == '__main__':
args = parse_args()
if args.list:
pass
#make_list(args)
else:
if os.path.isdir(args.prefix):
working_dir = args.prefix
else:
working_dir = os.path.dirname(args.prefix)
print('working_dir', working_dir)
prop = face_image.load_property(working_dir)
image_size = prop.image_size
print('image_size', image_size)
args.image_h = image_size[0]
args.image_w = image_size[1]
files = [
os.path.join(working_dir, fname)
for fname in os.listdir(working_dir)
if os.path.isfile(os.path.join(working_dir, fname))
]
count = 0
for fname in files:
if fname.startswith(args.prefix) and fname.endswith('.lst'):
print('Creating .rec file from', fname, 'in', working_dir)
count += 1
image_list = read_list(fname)
def train_net(args):
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))
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.rescale_threshold = 0
args.image_channel = 3
os.environ['BETA'] = str(args.beta)
data_dir_list = args.data_dir.split(',')
assert len(data_dir_list) == 1
data_dir = data_dir_list[0]
path_imgrec = None
path_imglist = None
prop = face_image.load_property(data_dir)
args.num_classes = prop.num_classes
#image_size = prop.image_size
image_size = [int(x) for x in args.image_size.split(',')]
assert len(image_size) == 2
assert image_size[0] == image_size[1]
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")
if args.loss_type == 1 and args.num_classes > 20000:
args.beta_freeze = 5000
args.gamma = 0.06
print('Called with argument:', args)
data_shape = (args.image_channel, image_size[0], image_size[1])
mean = None
begin_epoch = 0
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained) == 0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
if args.network[0] == 's':
data_shape_dict = {'data': (args.per_batch_size, ) + data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
else:
vec = args.pretrained.split(',')
print('loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(
vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
#label_name = 'softmax_label'
#label_shape = (args.batch_size,)
model = mx.mod.Module(
context=ctx,
symbol=sym,
)
val_dataiter = None
train_dataiter = 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,
color_jittering=args.color,
images_filter=args.images_filter,
)
metric1 = AccMetric()
eval_metrics = [mx.metric.create(metric1)]
if args.ce_loss:
metric2 = LossValueMetric()
eval_metrics.append(mx.metric.create(metric2))
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)
#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=base_lr,
momentum=base_mom,
wd=base_wd,
rescale_grad=_rescale)
som = 20
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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 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]
if len(args.lr_steps) == 0:
lr_steps = [40000, 60000, 80000]
if args.loss_type >= 1 and args.loss_type <= 7:
lr_steps = [100000, 140000, 160000]
p = 512.0 / args.batch_size
for l in range(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)
def _batch_callback(param):
#global global_step
global_step[0] += 1
mbatch = global_step[0]
for _lr in lr_steps:
if mbatch == args.beta_freeze + _lr:
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()
mx.model.save_checkpoint(prefix, msave, model.symbol, arg, aux)
print('[%d]Accuracy-Highest: %1.5f' % (mbatch, highest_acc[-1]))
if mbatch <= args.beta_freeze:
_beta = args.beta
else:
move = max(0, mbatch - args.beta_freeze)
_beta = max(
args.beta_min,
args.beta * math.pow(1 + args.gamma * move, -1.0 * args.power))
#print('beta', _beta)
os.environ['BETA'] = str(_beta)
if args.max_steps > 0 and mbatch > args.max_steps:
sys.exit(0)
epoch_cb = None
train_dataiter = mx.io.PrefetchingIter(train_dataiter)
model.fit(
train_dataiter,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
eval_data=val_dataiter,
eval_metric=eval_metrics,
kvstore='device',
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):
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))
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_list = args.data_dir.split(',')
assert len(data_dir_list)==1
data_dir = data_dir_list[0]
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_imglist = "/raid5data/dplearn/MS-Celeb-Aligned/lst2"
path_imgrec = os.path.join(data_dir, "train.rec")
assert args.images_per_identity>=2
assert args.triplet_bag_size%args.batch_size==0
print('Called with argument:', args)
data_shape = (args.image_channel,image_size[0],image_size[1])
mean = None
begin_epoch = 0
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained)==0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(',')
print('loading', vec)
sym, arg_params, aux_params = mx.model.load_checkpoint(vec[0], int(vec[1]))
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params, sym_embedding = sym)
data_extra = None
hard_mining = False
triplet_params = [args.triplet_bag_size, args.triplet_alpha, args.triplet_max_ap]
model = mx.mod.Module(
context = ctx,
symbol = sym,
#data_names = ('data',),
#label_names = None,
#label_names = ('softmax_label',),
)
label_shape = (args.batch_size,)
val_dataiter = None
train_dataiter = 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,
ctx_num = args.ctx_num,
images_per_identity = args.images_per_identity,
triplet_params = triplet_params,
mx_model = model,
)
_metric = LossValueMetric()
eval_metrics = [mx.metric.create(_metric)]
if args.network[0]=='r':
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)
_rescale = 1.0/args.ctx_num
if args.noise_sgd>0.0:
print('use noise sgd')
opt = NoiseSGD(scale = args.noise_sgd, learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
else:
opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
som = 2
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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 range(len(ver_list)):
acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(ver_list[i], model, args.batch_size, 10, None, label_shape)
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]
if len(args.lr_steps)==0:
lr_steps = [1000000000]
else:
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 _lr in lr_steps:
if mbatch==_lr:
opt.lr *= 0.1
print('lr change to', opt.lr)
break
_cb(param)
if mbatch%100==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
print("acc_list: ")
print(*acc_list)
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)
logger.info("score = " + str(score))
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()
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)
#epoch_cb = mx.callback.do_checkpoint(prefix, 1)
epoch_cb = None
model.fit(train_dataiter,
begin_epoch = begin_epoch,
num_epoch = end_epoch,
eval_data = val_dataiter,
eval_metric = eval_metrics,
kvstore = 'device',
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 main(args):
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))
ctx_num = len(ctx)
path_imgrec = os.path.join(args.input, 'train.rec')
path_imgidx = os.path.join(args.input, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
outf = open(os.path.join(args.input, 'c2c'), 'w')
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag > 0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
id2range = {}
seq_identity = range(int(header.label[0]), int(header.label[1]))
for identity in seq_identity:
s = imgrec.read_idx(identity)
header, _ = mx.recordio.unpack(s)
id2range[identity] = (int(header.label[0]), int(header.label[1]))
print('id2range', len(id2range))
prop = face_image.load_property(args.input)
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading', prefix, epoch)
model = mx.mod.Module.load(prefix, epoch, context=ctx)
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0],
image_size[1]))],
label_shapes=[('softmax_label', (args.batch_size, ))])
nrof_images = 0
nrof_removed = 0
idx = 1
id2label = {}
pp = 0
for _id, v in id2range.iteritems():
pp += 1
if pp % 100 == 0:
print('processing id', pp)
_list = range(*v)
ocontents = []
for i in range(len(_list)):
_idx = _list[i]
#print('_idx', _id, _idx)
s = imgrec.read_idx(_idx)
ocontents.append(s)
#continue
embeddings = None
headers = [None] * len(ocontents)
#print(len(ocontents))
ba = 0
while True:
bb = min(ba + args.batch_size, len(ocontents))
if ba >= bb:
break
_batch_size = bb - ba
_batch_size2 = max(_batch_size, ctx_num)
data = nd.zeros((_batch_size2, 3, image_size[0], image_size[1]))
label = nd.zeros((_batch_size2, ))
count = bb - ba
ii = 0
for i in range(ba, bb):
header, img = mx.recordio.unpack(ocontents[i])
headers[i] = header
img = mx.image.imdecode(img)
img = nd.transpose(img, axes=(2, 0, 1))
data[ii][:] = img
label[ii][:] = header.label[0]
ii += 1
while ii < _batch_size2:
data[ii][:] = data[0][:]
label[ii][:] = label[0][:]
ii += 1
db = mx.io.DataBatch(data=(data, ), label=(label, ))
model.forward(db, is_train=False)
net_out = model.get_outputs()
net_out = net_out[0].asnumpy()
if embeddings is None:
embeddings = np.zeros((len(ocontents), net_out.shape[1]))
embeddings[ba:bb, :] = net_out[0:_batch_size, :]
ba = bb
embeddings = sklearn.preprocessing.normalize(embeddings)
emb_mean = np.mean(embeddings, axis=0, keepdims=True)
emb_mean = sklearn.preprocessing.normalize(emb_mean)
sim = np.dot(embeddings, emb_mean.T)
#print(sim.shape)
sims = sim.flatten()
assert len(_list) == len(sims)
assert len(_list) == len(ocontents)
for i in range(len(ocontents)):
_sim = sims[i]
_idx = _list[i]
_header = headers[i]
#TODO
outf.write("%d,%f,%d\n" % (_idx, _sim, int(_header.label[1])))
outf.close()
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.rescale_threshold = 0
args.image_channel = 3
data_dir_list = args.data_dir.split(',')
assert len(data_dir_list) == 1
data_dir = data_dir_list[0]
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
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained) == 0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(',')
print('loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(
vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
if args.network[0] == 's':
data_shape_dict = {'data': (args.per_batch_size, ) + data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
#label_name = 'softmax_label'
#label_shape = (args.batch_size,)
model = mx.mod.Module(
context=ctx,
symbol=sym,
)
train_dataiter = 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,
)
val_rec = os.path.join(data_dir, "val.rec")
val_iter = None
if os.path.exists(val_rec):
val_iter = FaceImageIter(
batch_size=args.batch_size,
data_shape=data_shape,
path_imgrec=val_rec,
shuffle=False,
rand_mirror=False,
mean=mean,
)
eval_metrics = []
if USE_FR:
_metric = AccMetric(pred_idx=1)
eval_metrics.append(_metric)
if USE_GENDER:
_metric = AccMetric(pred_idx=2, name='gender')
eval_metrics.append(_metric)
elif USE_GENDER:
_metric = AccMetric(pred_idx=1, name='gender')
eval_metrics.append(_metric)
if USE_AGE:
_metric = MAEMetric()
eval_metrics.append(_metric)
_metric = CUMMetric()
eval_metrics.append(_metric)
if args.network[0] == 'r':
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)
_rescale = 1.0 / args.ctx_num
#opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
opt = optimizer.Nadam(learning_rate=base_lr,
wd=base_wd,
rescale_grad=_rescale)
som = 20
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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], 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
def val_test():
_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 i, eval_batch in enumerate(val_iter):
model.forward(eval_batch, is_train=False)
model.update_metric(val_metric, eval_batch.label)
model.update_metric(val_metric2, eval_batch.label)
_value = val_metric.get_name_value()[0][1]
print('MAE: %f' % (_value))
_value = val_metric2.get_name_value()[0][1]
print('CUM: %f' % (_value))
highest_acc = [0.0, 0.0] #lfw and target
#for i in xrange(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
if len(args.lr_steps) == 0:
lr_steps = [40000, 60000, 80000]
if args.loss_type >= 1 and args.loss_type <= 7:
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)
def _batch_callback(param):
#global global_step
global_step[0] += 1
mbatch = global_step[0]
for _lr in lr_steps:
if mbatch == _lr:
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:
if val_iter is not None:
val_test()
acc_list = ver_test(mbatch)
save_step[0] += 1
msave = save_step[0]
do_save = 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
if acc_list[-1] >= highest_acc[-1]:
highest_acc[-1] = acc_list[-1]
if lfw_score >= 0.99:
do_save = True
if args.ckpt == 0:
do_save = False
elif args.ckpt > 1:
do_save = True
if do_save:
print('saving', 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)
epoch_cb = None
model.fit(
train_dataiter,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
eval_data=None,
eval_metric=eval_metrics,
kvstore='device',
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)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='do verification')
# general
parser.add_argument('--data-dir', default='', help='')
parser.add_argument('--model', default='../model/softmax,50', help='path to load model.')
parser.add_argument('--target', default='lfw,cfp_ff,cfp_fp,agedb_30', help='test targets.')
parser.add_argument('--gpu', default=0, type=int, help='gpu id')
parser.add_argument('--batch-size', default=32, type=int, help='')
parser.add_argument('--max', default='', type=str, help='')
parser.add_argument('--mode', default=0, type=int, help='')
parser.add_argument('--nfolds', default=10, type=int, help='')
args = parser.parse_args()
prop = face_image.load_property(args.data_dir)
image_size = prop.image_size
print('image_size', image_size)
ctx = mx.gpu(args.gpu)
nets = []
vec = args.model.split(',')
prefix = args.model.split(',')[0]
epochs = []
if len(vec)==1:
pdir = os.path.dirname(prefix)
for fname in os.listdir(pdir):
if not fname.endswith('.params'):
continue
_file = os.path.join(pdir, fname)
if _file.startswith(prefix):
epoch = int(fname.split('.')[0].split('-')[1])
def train_net(args):
ctx = []
cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip() #0,使用第一块GPU
if len(cvd) > 0:
for i in xrange(len(cvd.split(','))):
ctx.append(mx.gpu(i)) #讲GPU context添加到ctx,ctx = [gpu(0)]
if len(ctx) == 0:
ctx = [mx.cpu()]
print('use cpu')
else:
print('gpu num:', len(ctx)) #使用了gpu
prefix = args.prefix #../model-r100
prefix_dir = os.path.dirname(prefix) #..
if not os.path.exists(prefix_dir): #未执行
os.makedirs(prefix_dir)
end_epoch = args.end_epoch #100 000
args.ctx_num = len(ctx)
args.num_layers = int(args.network[1:])
print('num_layers', args.num_layers) #100
if args.per_batch_size == 0:
args.per_batch_size = 128
args.batch_size = args.per_batch_size * args.ctx_num #10
args.rescale_threshold = 0
args.image_channel = 3
os.environ['BETA'] = str(args.beta) #1000.0,参见Arcface公式(6)，退火训练的lambda
data_dir_list = args.data_dir.split(',')
print('data_dir_list: ', data_dir_list)
data_dir = data_dir_list[0]
# 加载数据集属性
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)
print('num_classes: ', args.num_classes)
path_imgrec = os.path.join(data_dir, "train.rec")
if args.loss_type == 1 and args.num_classes > 20000: #sphereface
args.beta_freeze = 5000
args.gamma = 0.06
print('***Called with argument:', args)
data_shape = (args.image_channel, image_size[0], image_size[1]
) #(3L,112L,112L)
mean = None
begin_epoch = 0
base_lr = args.lr #0.1
base_wd = args.wd #weight decay = 0.0005
base_mom = args.mom #动量:0.9
if len(args.pretrained) == 0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(
',') #['../models/model-r50-am-lfw/model', '0000']
print('***loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(
vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
# print(sym[1])
# mx.viz.plot_network(sym[1]).view() #可视化
# sys.exit()
if args.network[0] == 's': # spherenet
data_shape_dict = {'data': (args.per_batch_size, ) + data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
#label_name = 'softmax_label'
#label_shape = (args.batch_size,)
model = mx.mod.Module(
context=ctx,
symbol=sym,
)
# print(args.batch_size)
# print(data_shape)
# print(path_imgrec)
# print(args.rand_mirror)
# print(mean)
# print(args.cutoff)
# sys.exit()
train_dataiter = FaceImageIter(
batch_size=args.batch_size,
data_shape=data_shape, #(3L,112L,112L)
path_imgrec=path_imgrec, # train.rec
shuffle=True,
rand_mirror=args.rand_mirror, # 1
mean=mean,
cutoff=args.cutoff, # 0
)
if args.loss_type < 10:
_metric = AccMetric()
else:
_metric = LossValueMetric()
# 创建一个评价指标
eval_metrics = [mx.metric.create(_metric)]
if args.network[0] == 'r' or args.network[0] == 'y':
initializer = mx.init.Xavier(rnd_type='gaussian',
factor_type="out",
magnitude=2) #resnet style mobilefacenet
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)
_rescale = 1.0 / args.ctx_num
opt = optimizer.SGD(learning_rate=base_lr,
momentum=base_mom,
wd=base_wd,
rescale_grad=_rescale) #多卡训练的话，rescale_grad将总的结果分开
som = 64
# 回调函数，用来阶段性显示训练速度和准确率
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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], 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 xrange(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
if len(args.lr_steps) == 0:
lr_steps = [30000, 40000, 50000]
if args.loss_type >= 1 and args.loss_type <= 7:
lr_steps = [100000, 140000, 160000]
# 单GPU，去掉p
# p = 512.0/args.batch_size
for l in xrange(len(lr_steps)):
# lr_steps[l] = int(lr_steps[l]*p)
lr_steps[l] = int(lr_steps[l])
else:
lr_steps = [int(x) for x in args.lr_steps.split(',')]
print('lr_steps', lr_steps)
def _batch_callback(param):
#global global_step
mbatch = global_step[0]
global_step[0] += 1
for _lr in lr_steps:
if mbatch == args.beta_freeze + _lr:
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)
print(acc_list)
save_step[0] += 1
msave = save_step[0]
do_save = 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:
if lfw_score >= 0.99:
do_save = True
if acc_list[-1] >= highest_acc[-1]:
highest_acc[-1] = acc_list[-1]
# if lfw_score>=0.99: #LFW测试大于0.99时,保存模型
if lfw_score >= 0.99: #LFW测试大于0.99时,保存模型
do_save = True
if args.ckpt == 0:
do_save = False
elif args.ckpt > 1:
do_save = True
if do_save:
print('saving', 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 mbatch <= args.beta_freeze:
_beta = args.beta
else:
move = max(0, mbatch - args.beta_freeze)
_beta = max(
args.beta_min,
args.beta * math.pow(1 + args.gamma * move, -1.0 * args.power))
#print('beta', _beta)
os.environ['BETA'] = str(_beta)
if args.max_steps > 0 and mbatch > args.max_steps:
sys.exit(0)
epoch_cb = None
train_dataiter = mx.io.PrefetchingIter(train_dataiter)
model.fit(
train_data=train_dataiter,
begin_epoch=begin_epoch,
num_epoch=end_epoch,
eval_data=None,
eval_metric=eval_metrics,
kvstore='device',
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 main(args):
include_datasets = args.include.split(',')
prop = face_image.load_property(include_datasets[0])
image_size = prop.image_size
print('image_size', image_size)
model = None
if len(args.model)>0:
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))
args.ctx_num = len(ctx)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading',prefix, epoch)
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
#arg_params, aux_params = ch_dev(arg_params, aux_params, ctx)
all_layers = sym.get_internals()
sym = all_layers['fc1_output']
#model = mx.mod.Module.load(prefix, epoch, context = ctx)
#model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))])
model.set_params(arg_params, aux_params)
else:
assert args.param1==0.0
rec_list = []
for ds in include_datasets:
path_imgrec = os.path.join(ds, 'train.rec')
path_imgidx = os.path.join(ds, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
rec_list.append(imgrec)
id_list_map = {}
all_id_list = []
test_limit = 0
for ds_id in xrange(len(rec_list)):
id_list = []
imgrec = rec_list[ds_id]
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag>0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
id2range = {}
seq_identity = range(int(header.label[0]), int(header.label[1]))
pp=0
for identity in seq_identity:
pp+=1
if pp%10==0:
print('processing id', pp)
if model is not None:
embedding = get_embedding(args, imgrec, identity, image_size, model)
else:
embedding = None
#print(embedding.shape)
id_list.append( [ds_id, identity, embedding] )
if test_limit>0 and pp>=test_limit:
break
id_list_map[ds_id] = id_list
if ds_id==0 or model is None:
all_id_list += id_list
print(ds_id, len(id_list))
else:
X = []
for id_item in all_id_list:
X.append(id_item[2])
X = np.array(X)
for i in xrange(len(id_list)):
id_item = id_list[i]
y = id_item[2]
sim = np.dot(X, y.T)
idx = np.where(sim>=args.param1)[0]
if len(idx)>0:
continue
all_id_list.append(id_item)
print(ds_id, len(id_list), len(all_id_list))
if len(args.exclude)>0:
if os.path.isdir(args.exclude):
_path_imgrec = os.path.join(args.exclude, 'train.rec')
_path_imgidx = os.path.join(args.exclude, 'train.idx')
_imgrec = mx.recordio.MXIndexedRecordIO(_path_imgidx, _path_imgrec, 'r') # pylint: disable=redefined-variable-type
_ds_id = len(rec_list)
_id_list = []
s = _imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag>0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
seq_identity = range(int(header.label[0]), int(header.label[1]))
pp=0
for identity in seq_identity:
pp+=1
if pp%10==0:
print('processing ex id', pp)
embedding = get_embedding(args, _imgrec, identity, image_size, model)
#print(embedding.shape)
_id_list.append( (_ds_id, identity, embedding) )
if test_limit>0 and pp>=test_limit:
break
else:
_id_list = []
data_set = verification.load_bin(args.exclude, image_size)[0][0]
print(data_set.shape)
data = nd.zeros( (1,3,image_size[0], image_size[1]))
for i in xrange(data_set.shape[0]):
data[0] = data_set[i]
db = mx.io.DataBatch(data=(data,))
model.forward(db, is_train=False)
net_out = model.get_outputs()
embedding = net_out[0].asnumpy().flatten()
_norm=np.linalg.norm(embedding)
embedding /= _norm
_id_list.append( (i, i, embedding) )
#X = []
#for id_item in all_id_list:
# X.append(id_item[2])
#X = np.array(X)
#param1 = 0.3
#while param1<=1.01:
# emap = {}
# for id_item in _id_list:
# y = id_item[2]
# sim = np.dot(X, y.T)
# #print(sim.shape)
# #print(sim)
# idx = np.where(sim>=param1)[0]
# for j in idx:
# emap[j] = 1
# exclude_removed = len(emap)
# print(param1, exclude_removed)
# param1+=0.05
X = []
for id_item in all_id_list:
X.append(id_item[2])
X = np.array(X)
emap = {}
for id_item in _id_list:
y = id_item[2]
sim = np.dot(X, y.T)
idx = np.where(sim>=args.param2)[0]
for j in idx:
emap[j] = 1
all_id_list[j][1] = -1
print('exclude', len(emap))
if args.test>0:
return
if not os.path.exists(args.output):
os.makedirs(args.output)
writer = mx.recordio.MXIndexedRecordIO(os.path.join(args.output, 'train.idx'), os.path.join(args.output, 'train.rec'), 'w')
idx = 1
identities = []
nlabel = -1
for id_item in all_id_list:
if id_item[1]<0:
continue
nlabel+=1
ds_id = id_item[0]
imgrec = rec_list[ds_id]
id = id_item[1]
s = imgrec.read_idx(id)
header, _ = mx.recordio.unpack(s)
a, b = int(header.label[0]), int(header.label[1])
identities.append( (idx, idx+b-a) )
for _idx in xrange(a,b):
s = imgrec.read_idx(_idx)
_header, _content = mx.recordio.unpack(s)
nheader = mx.recordio.IRHeader(0, nlabel, idx, 0)
s = mx.recordio.pack(nheader, _content)
writer.write_idx(idx, s)
idx+=1
id_idx = idx
for id_label in identities:
_header = mx.recordio.IRHeader(1, id_label, idx, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(idx, s)
idx+=1
_header = mx.recordio.IRHeader(1, (id_idx, idx), 0, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(0, s)
with open(os.path.join(args.output, 'property'), 'w') as f:
f.write("%d,%d,%d"%(len(identities), image_size[0], image_size[1]))
def train_net(args):
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))
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 range(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 range(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 range(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)))
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.rescale_threshold = 0
args.image_channel = 3
data_dir_list = args.data_dir.split(',')
assert len(data_dir_list)==1
data_dir = data_dir_list[0]
path_imgrec = None
path_imglist = None
args.num_classes = 0
image_size = (112,112)
if os.path.exists(os.path.join(data_dir, 'property')):
prop = face_image.load_property(data_dir)
args.num_classes = prop.num_classes
image_size = prop.image_size
assert(args.num_classes>0)
print('num_classes', args.num_classes)
args.image_h = image_size[0]
args.image_w = image_size[1]
print('image_size', image_size)
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
base_lr = args.lr
base_wd = args.wd
base_mom = args.mom
if len(args.pretrained)==0:
arg_params = None
aux_params = None
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
else:
vec = args.pretrained.split(',')
print('loading', vec)
_, arg_params, aux_params = mx.model.load_checkpoint(vec[0], int(vec[1]))
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
if args.network[0]=='s':
data_shape_dict = {'data' : (args.per_batch_size,)+data_shape}
spherenet.init_weights(sym, data_shape_dict, args.num_layers)
#label_name = 'softmax_label'
#label_shape = (args.batch_size,)
model = mx.mod.Module(
context = ctx,
symbol = sym,
)
train_dataiter = 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,
)
val_rec = os.path.join(data_dir, "val.rec")
val_iter = None
if os.path.exists(val_rec):
val_iter = FaceImageIter(
batch_size = args.batch_size,
data_shape = data_shape,
path_imgrec = val_rec,
shuffle = False,
rand_mirror = False,
mean = mean,
)
eval_metrics = []
if USE_FR:
_metric = AccMetric(pred_idx=1, label_idx=0)
eval_metrics.append(_metric)
if USE_GENDER:
_metric = AccMetric(pred_idx=2, label_idx=1, name='gender')
eval_metrics.append(_metric)
elif USE_GENDER:
_metric = AccMetric(pred_idx=1, label_idx=1, name='gender')
eval_metrics.append(_metric)
if USE_AGE:
_metric = MAEMetric()
eval_metrics.append(_metric)
_metric = CUMMetric()
eval_metrics.append(_metric)
if args.network[0]=='r':
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)
_rescale = 1.0/args.ctx_num
opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=_rescale)
#opt = optimizer.Nadam(learning_rate=base_lr, wd=base_wd, rescale_grad=_rescale)
som = 20
_cb = mx.callback.Speedometer(args.batch_size, som)
ver_list = []
ver_name_list = []
for name in args.target.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], 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
def val_test():
_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 i, eval_batch in enumerate(val_iter):
model.forward(eval_batch, is_train=False)
model.update_metric(val_metric, eval_batch.label)
model.update_metric(val_metric2, eval_batch.label)
_value = val_metric.get_name_value()[0][1]
print('MAE: %f'%(_value))
_value = val_metric2.get_name_value()[0][1]
print('CUM: %f'%(_value))
highest_acc = [0.0, 0.0] #lfw and target
#for i in xrange(len(ver_list)):
# highest_acc.append(0.0)
global_step = [0]
save_step = [0]
if len(args.lr_steps)==0:
lr_steps = [40000, 60000, 80000]
if args.loss_type>=1 and args.loss_type<=7:
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)
def _batch_callback(param):
#global global_step
global_step[0]+=1
mbatch = global_step[0]
for _lr in lr_steps:
if mbatch==_lr:
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:
if val_iter is not None:
val_test()
acc_list = ver_test(mbatch)
save_step[0]+=1
msave = save_step[0]
do_save = 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
if acc_list[-1]>=highest_acc[-1]:
highest_acc[-1] = acc_list[-1]
if lfw_score>=0.99:
do_save = True
if args.ckpt==0:
do_save = False
elif args.ckpt>1:
do_save = True
if do_save:
print('saving', 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)
epoch_cb = None
model.fit(train_dataiter,
begin_epoch = begin_epoch,
num_epoch = end_epoch,
eval_data = None,
eval_metric = eval_metrics,
kvstore = 'device',
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 main(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))
args.ctx_num = len(ctx)
include_datasets = args.include.split(',')
prop = face_image.load_property(include_datasets[0])
image_size = prop.image_size
print('image_size', image_size)
vec = args.model.split(',')
prefix = vec[0]
epoch = int(vec[1])
print('loading', prefix, epoch)
model = mx.mod.Module.load(prefix, epoch, context=ctx)
model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0],
image_size[1]))],
label_shapes=[('softmax_label', (args.batch_size, ))])
rec_list = []
for ds in include_datasets:
path_imgrec = os.path.join(ds, 'train.rec')
path_imgidx = os.path.join(ds, 'train.idx')
imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
rec_list.append(imgrec)
id_list_map = {}
all_id_list = []
test_limit = 0
for ds_id in xrange(len(rec_list)):
id_list = []
imgrec = rec_list[ds_id]
s = imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag > 0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
id2range = {}
seq_identity = range(int(header.label[0]), int(header.label[1]))
pp = 0
for identity in seq_identity:
pp += 1
if pp % 10 == 0:
print('processing id', pp)
embedding = get_embedding(args, imgrec, identity, image_size,
model)
#print(embedding.shape)
id_list.append([ds_id, identity, embedding])
if test_limit > 0 and pp >= test_limit:
break
id_list_map[ds_id] = id_list
if ds_id == 0:
all_id_list += id_list
print(ds_id, len(id_list))
else:
X = []
for id_item in all_id_list:
X.append(id_item[2])
X = np.array(X)
for i in xrange(len(id_list)):
id_item = id_list[i]
y = id_item[2]
sim = np.dot(X, y.T)
idx = np.where(sim >= args.param1)[0]
if len(idx) > 0:
continue
all_id_list.append(id_item)
print(ds_id, len(id_list), len(all_id_list))
if len(args.exclude) > 0:
_path_imgrec = os.path.join(args.exclude, 'train.rec')
_path_imgidx = os.path.join(args.exclude, 'train.idx')
_imgrec = mx.recordio.MXIndexedRecordIO(_path_imgidx, _path_imgrec,
'r') # pylint: disable=redefined-variable-type
_ds_id = len(rec_list)
_id_list = []
s = _imgrec.read_idx(0)
header, _ = mx.recordio.unpack(s)
assert header.flag > 0
print('header0 label', header.label)
header0 = (int(header.label[0]), int(header.label[1]))
#assert(header.flag==1)
imgidx = range(1, int(header.label[0]))
seq_identity = range(int(header.label[0]), int(header.label[1]))
pp = 0
for identity in seq_identity:
pp += 1
if pp % 10 == 0:
print('processing ex id', pp)
embedding = get_embedding(args, _imgrec, identity, image_size,
model)
#print(embedding.shape)
_id_list.append((_ds_id, identity, embedding))
if test_limit > 0 and pp >= test_limit:
break
#X = []
#for id_item in all_id_list:
# X.append(id_item[2])
#X = np.array(X)
#param1 = 0.3
#while param1<=1.01:
# emap = {}
# for id_item in _id_list:
# y = id_item[2]
# sim = np.dot(X, y.T)
# #print(sim.shape)
# #print(sim)
# idx = np.where(sim>=param1)[0]
# for j in idx:
# emap[j] = 1
# exclude_removed = len(emap)
# print(param1, exclude_removed)
# param1+=0.05
X = []
for id_item in all_id_list:
X.append(id_item[2])
X = np.array(X)
emap = {}
for id_item in _id_list:
y = id_item[2]
sim = np.dot(X, y.T)
idx = np.where(sim >= args.param2)[0]
for j in idx:
emap[j] = 1
all_id_list[j][1] = -1
print('exclude', len(emap))
if args.test > 0:
return
if not os.path.exists(args.output):
os.makedirs(args.output)
writer = mx.recordio.MXIndexedRecordIO(
os.path.join(args.output, 'train.idx'),
os.path.join(args.output, 'train.rec'), 'w')
idx = 1
identities = []
nlabel = -1
for id_item in all_id_list:
if id_item[1] < 0:
continue
nlabel += 1
ds_id = id_item[0]
imgrec = rec_list[ds_id]
id = id_item[1]
s = imgrec.read_idx(id)
header, _ = mx.recordio.unpack(s)
a, b = int(header.label[0]), int(header.label[1])
identities.append((idx, idx + b - a))
for _idx in xrange(a, b):
s = imgrec.read_idx(_idx)
_header, _content = mx.recordio.unpack(s)
nheader = mx.recordio.IRHeader(0, nlabel, idx, 0)
s = mx.recordio.pack(nheader, _content)
writer.write_idx(idx, s)
idx += 1
id_idx = idx
for id_label in identities:
_header = mx.recordio.IRHeader(1, id_label, idx, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(idx, s)
idx += 1
_header = mx.recordio.IRHeader(1, (id_idx, idx), 0, 0)
s = mx.recordio.pack(_header, '')
writer.write_idx(0, s)
with open(os.path.join(args.output, 'property'), 'w') as f:
f.write("%d,%d,%d" % (len(identities), image_size[0], image_size[1]))
