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 forward(self, x):
# x: (n,c,w)
# X_ = dot(W1, input)
#import pdb
#pdb.set_trace()
x = F.transpose(x, axes=(0, 2, 1)) # (nwc) -> (ncw)
X_ = F.batch_dot(self.w1.data(ctx), x) # (n,c,w) -> (n,c,w)
# E = dot(X_, W)
E = F.batch_dot(X_, self.w.data(ctx)) # (n,c,w) -> (n,c,w)
attn_weights = F.softmax(E, axis=2) # (n, c, w)
attn_applied = F.elemwise_mul(attn_weights, X_) #(n,c,w)
output = self.c.data(ctx) * (attn_applied) + (
1 - self.c.data(ctx)) * X_ # (n,c,w)
output = F.batch_dot(output, self.w2.data(ctx)) + self.b.data(
ctx) # (n, c,w)
output = F.transpose(output, axes=(0, 2, 1)) # (ncw) -> (nwc)
return output
def fn(heads, relations, tails, num_chunks, chunk_size,
neg_sample_size):
hidden_dim = heads.shape[1]
tails = tails.reshape(num_chunks, neg_sample_size, hidden_dim)
tails = nd.transpose(tails, axes=(0, 2, 1))
tmp = (heads * relations).reshape(num_chunks, chunk_size,
hidden_dim)
return nd.linalg_gemm2(tmp, tails)
def post_process_data(image_arr):
"""
[height, width, channel] ---> [channel, height, width]
:param image_arr:
:return:
"""
return nd.transpose(image_arr, axes=(2, 0, 1))
pass
def AUG_Transform_data_label(data, label):
data = nd.array(data)
data = mx.image.imresize(data, 112, 112)
light_aug = mx.gluon.data.vision.transforms.RandomLighting(0.3)
bright_aug = mx.gluon.data.vision.transforms.RandomBrightness(0.3)
flip_aug = mx.gluon.data.vision.transforms.RandomFlipLeftRight()
data = bright_aug(light_aug(flip_aug(data)))
return nd.transpose(data.astype('float32'), (2, 0, 1)) / 255.0, nd.array(
[label]).asscalar().astype('float32')
def transform_mnist(data, label): # 处理数据格式
if resize:
n = data.shape[0]
new_data = nd.zeros((n, resize, resize, data.shape[3]))
for i in range(n):
new_data[i] = image.imresize(data[i], resize, resize)
data = new_data
return nd.transpose(data.astype('float32'),
(0, 3, 1, 2)) / 255, label.astype('float32')
def PredictTrans(predict):
colormap = ndarray.array(MICCAI_colormap, ctx=mx.gpu(),
dtype='uint8') # voc_colormap
label = colormap[predict[:, :, :]]
label = ndarray.transpose(label, (0, 3, 1, 2))
label = label.astype(('float32'))
return label
def transform_mnist(data, label):
# transform a batch of examples
if self.__resize:
# data 默认 (28, 28, 1)
# data imresize 后 (224, 224, 1)
data = image.imresize(data, self.__resize, self.__resize)
# change data from height x weight x channel to channel x height x weight
return nd.transpose(data.astype("float32"),
(2, 0, 1)) / 255, label.astype("float32")
def transform(data, target_wd, target_ht):
# resize to target_wd * target_ht
data = mx.image.imresize(data, target_wd, target_ht)
# transpose from (target_wd, target_ht, 3)
# to (3, target_wd, target_ht)
data = nd.transpose(data, (2, 0, 1))
# normalize to [-1, 1]
data = data.astype(np.float32) / 127.5 - 1
return data.reshape((1,) + data.shape)
def transform_mask(mask):
# Convert types
mask = mask.astype('float32') / 255
# Convert mask to binary
mask = (mask > 0.4).astype('float32')
# Reshape the tensors so the order is now (channels, w, h)
return nd.transpose(mask, (2, 0, 1))
def _rearrange(raw, F, upscale_factor):
# (N, C * r^2, H, W) -> (N, C, r^2, H, W)
splitted = F.reshape(raw, shape=(0, -4, -1, upscale_factor**2, 0, 0))
# (N, C, r^2, H, W) -> (N, C, r, r, H, W)
unflatten = F.reshape(splitted, shape=(0, 0, -4, upscale_factor, upscale_factor, 0, 0))
# (N, C, r, r, H, W) -> (N, C, H, r, W, r)
swapped = F.transpose(unflatten, axes=(0, 1, 4, 2, 5, 3))
# (N, C, H, r, W, r) -> (N, C, H*r, W*r)
return F.reshape(swapped, shape=(0, 0, -3, -3))
def resize_img(data, resize=None):
if resize:
n = data.shape[0]
new_data = nd.zeros((n, resize, resize, data.shape[3]))
for i in range(n):
new_data[i] = image.imresize(data[i].as_in_context(mx.cpu()),
resize, resize)
data = new_data
return nd.transpose(data.astype('float32'), (0, 3, 1, 2)) / 255
def _rearrange(raw, F, upscale_factor):
# (N, C * r^2, H, W) -> (N, C, r^2, H, W)
splitted = F.reshape(raw, shape=(0, -4, -1, upscale_factor**2, 0, 0))
# (N, C, r^2, H, W) -> (N, C, r, r, H, W)
unflatten = F.reshape(splitted, shape=(0, 0, -4, upscale_factor, upscale_factor, 0, 0))
# (N, C, r, r, H, W) -> (N, C, H, r, W, r)
swapped = F.transpose(unflatten, axes=(0, 1, 4, 2, 5, 3))
# (N, C, H, r, W, r) -> (N, C, H*r, W*r)
return F.reshape(swapped, shape=(0, 0, -3, -3))
def batched_l2_dist(a, b):
a_squared = nd.power(nd.norm(a, axis=-1), 2)
b_squared = nd.power(nd.norm(b, axis=-1), 2)
squared_res = nd.add(nd.linalg_gemm(
a, nd.transpose(b, axes=(0, 2, 1)), nd.broadcast_axes(nd.expand_dims(b_squared, axis=-2), axis=1, size=a.shape[1]), alpha=-2
), nd.expand_dims(a_squared, axis=-1))
res = nd.sqrt(nd.clip(squared_res, 1e-30, np.finfo(np.float32).max))
return res
def predict2ndimg(predict):
result = ndarray.argmax(predict, axis=1)
colormap = ndarray.array(MICCAI_colormap, ctx=mx.gpu(),
dtype='uint8') # voc_colormap
ndimg = colormap[result[:, :, :]]
ndimg = ndarray.transpose(ndimg, (0, 3, 1, 2))
ndimg = ndimg.astype(('float32'))
return ndimg
def transform(data, label):
# data: sample x height x width x channel
# label: sample
data = data.astype('float32')
if augs is not None:
# apply to each sample one-by-one and then stack
data = nd.stack(*[apply_aug_list(d, augs) for d in data])
data = nd.transpose(data, (0, 3, 1, 2))
return data, label.astype('float32')
def load_test_images(fnames,
lbl,
batch_size,
img_wd,
img_ht,
ctx,
noisevar=0.2,
is_reversed=False,
bw=1):
img_in_list = []
img_out_list = []
#shuffle(fnames)
for img in fnames:
img_arr = mx.image.imread(img, flag=bw).astype(np.float32) / 127.5 - 1
img_arr = mx.image.imresize(img_arr, img_wd, img_ht)
# Crop input and output images
croppedimg = mx.image.fixed_crop(img_arr, 0, 0, img_wd, img_ht)
if noisevar > 0:
img_arr_in, img_arr_out = [
croppedimg + mx.random.normal(0, noisevar, croppedimg.shape),
croppedimg
]
else:
img_arr_in, img_arr_out = [croppedimg, croppedimg]
img_arr_in, img_arr_out = [
nd.transpose(img_arr_in, (2, 0, 1)),
nd.transpose(img_arr_out, (2, 0, 1))
]
img_arr_in, img_arr_out = [
img_arr_in.reshape((1, ) + img_arr_in.shape),
img_arr_out.reshape((1, ) + img_arr_out.shape)
]
img_in_list.append(img_arr_out if is_reversed else img_arr_in)
img_out_list.append(img_arr_in if is_reversed else img_arr_out)
tempdata = [
nd.concat(*img_in_list, dim=0),
nd.concat(*img_out_list, dim=0)
]
templbl = mx.nd.array(lbl)
itertest = mx.io.NDArrayIter(data=tempdata,
label=templbl,
batch_size=batch_size)
return itertest
def get_embedding(args, imgrec, id, image_size, model):
s = imgrec.read_idx(id)
header, _ = mx.recordio.unpack(s)
ocontents = []
# print('** get_embedding:', int(header.label[0]), int(header.label[1]))
# for idx in range(int(imgrec.keys[0]), int(imgrec.keys[-1])):
# for idx in range(int(id), int(id)+1):
upper_image_index = min(
args.max_image_per_identity_in_embedding + int(header.label[0]),
int(header.label[1]))
for idx in xrange(int(header.label[0]), upper_image_index):
# print('** get_embedding idx', idx)
s = imgrec.read_idx(idx)
ocontents.append(s)
embeddings = None
# print('len(ocontents)', 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, args.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])
#print(header.label.shape, header.label)
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,))
db = mx.io.DataBatch(data=(data, ))
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)
# print('get_embedding:embeddings.shape:', embeddings.shape)
embedding = np.mean(embeddings, axis=0, keepdims=True)
# print('get_embedding:embedding.shape:', embedding.shape)
embedding = sklearn.preprocessing.normalize(embedding).flatten()
# print('get_embedding:embedding.shape:', embedding.shape)
# print('get_embedding:', embedding)
return embedding
def hybrid_forward(self, F, x): # x.shape=(N, 1, C, 2)
# (n,1,c,2)->(n,c/k,k,2)
concat_step1 = x.reshape(
(x.shape[0], x.shape[2] / self._group_ratio, self._group_ratio, 2))
# (n,c/k,k,2)->(n,c/k,2,k)
concat_step2 = nd.transpose(concat_step1, (0, 1, 3, 2))
# (n,c/k,2,k)->(n, 2*c/k, k)
concat_out = concat_step2.reshape(
(concat_step2.shape[0], 1, -1, concat_step2.shape[-1]))
return concat_out
def transform_mnist(data, label):
# transform a batch of examples
if resize:
n = data.shape[0]
new_data = nd.zeros((n, resize, resize, data.shape[3]))
for i in range(n):
new_data[i] = image.imresize(data[i], resize, resize)
data = new_data
# change data from batch x height x weight x channel to batch x channel x height x weight
return nd.transpose(data.astype('float32'), (0,3,1,2))/255, label.astype('float32')
def transform_mnist(data, label):
# transform a batch of examples
if resize:
n = data.shape[0]
new_data = nd.zeros((n, resize, resize, data.shape[3]))
for i in range(n):
new_data[i] = image.imresize(data[i], resize, resize)
data = new_data
# change data from batch x height x weight x channel to batch x channel x height x weight
return nd.transpose(data.astype('float32'), (0,3,1,2))/255, label.astype('float32')
def show_img(data_iter):
while True:
try:
x = data_iter.next().data[0]
x = nd.transpose(x, (0, 2, 3, 1))
x = x[0].clip(0, 255) / 255
plt.imshow(x.asnumpy())
plt.show()
except StopIteration:
break
def load_data_fashion_mnist(batch_size, resize=None):
def transform(data, label):
if resize:
n = data.shape[0]
new_data = nd.zeros((n, resize, resize, data.shape[3]))
for i in range(n):
new_data[i] = image.imresize(data[i], resize, resize)
data = new_data
return nd.transpose(data.astype('float32'), (0,3,1,2))/255, label.astype('float32')
def RNN(epoch = 100 , batch_size=100, save_period=100 , load_period=100 ,learning_rate= 0.1, ctx=mx.gpu(0)):
train_data , test_data = FashionMNIST(batch_size)
#network parameter
time_step = 28
num_inputs = 28
num_hidden = 200
num_outputs = 10
path = "weights/FashionMNIST_RNNweights-{}.params".format(load_period)
model=RNNCell(num_hidden,num_outputs)
model.hybridize()
# weight initialization
if os.path.exists(path):
print("loading weights")
model.load_params(filename=path ,ctx=ctx) # weights load
else:
print("initializing weights")
model.collect_params().initialize(mx.init.Normal(sigma=0.01),ctx=ctx) # weights initialization
trainer = gluon.Trainer(model.collect_params(), "adam", {"learning_rate": learning_rate})
for i in tqdm(range(1,epoch+1,1)):
for data,label in train_data:
states = [nd.zeros(shape=(data.shape[0], num_hidden), ctx=ctx)]
data=data.as_in_context(ctx)
data = data.reshape(shape=(-1,time_step,num_inputs))
data= nd.transpose(data=data,axes=(1,0,2))
label = label.as_in_context(ctx)
with autograd.record():
for j in range(time_step):
outputs , states = model(data[j],states) #outputs => (batch size, 10)
loss = gluon.loss.SoftmaxCrossEntropyLoss()(outputs,label) # (batch_size,)
loss.backward()
trainer.step(batch_size)
cost = nd.mean(loss).asscalar()
test_accuracy = evaluate_accuracy(test_data, time_step, num_inputs, num_hidden, model, ctx)
print(" epoch : {} , last batch cost : {}".format(i,cost))
print("Test_acc : {0:0.3f}%".format(test_accuracy * 100))
#weight_save
if i % save_period==0:
if not os.path.exists("weights"):
os.makedirs("weights")
print("saving weights")
model.save_params("weights/FashionMNIST_RNNweights-{}.params".format(i))
test_accuracy = evaluate_accuracy(test_data, time_step, num_inputs, num_hidden, model, ctx)
print("Test_acc : {0:0.3f}%".format(test_accuracy * 100))
def forward(self,
encoder_output: nd.NDArray,
label=None,
label_lengths=None):
no_label = label is None or label_lengths is None
encoder_output = nd.transpose(encoder_output, (0, 2, 3, 1))
encoder_output = encoder_output.reshape(
(encoder_output.shape[0], -1, encoder_output.shape[3]))
batch_max_len = self.max_len if no_label else int(
label_lengths.max().asscalar()) - 1
# Initialize hidden states
encoder_output_mean = encoder_output.mean(axis=1)
h = self.init_h(encoder_output_mean)
c = self.init_c(encoder_output_mean)
# Two tensors to store outputs
predictions = []
alphas = []
if not no_label:
label_embedded = self.embedding(label)
else:
bs = encoder_output.shape[0]
x_t = self.embedding(
nd.zeros(shape=(bs, ), ctx=encoder_output.context))
for t in range(batch_max_len):
if not no_label:
x_t = label_embedded[:, t]
if self._use_current_state:
_, [h, c] = self.lstm_cell(x_t, [h, c])
if self._use_adaptive_attention:
atten_weights, alpha = self.attention(
encoder_output, h, x_t, c)
else:
atten_weights, alpha = self.attention(encoder_output, h)
atten_weights = self.f_beta(h).sigmoid() * atten_weights
inputs = nd.concat(atten_weights, h, dim=1)
preds = self.out(self.dropout(inputs))
else:
atten_weights, alpha = self.attention(encoder_output, h)
atten_weights = nd.sigmoid(self.f_beta(h)) * atten_weights
inputs = nd.concat(x_t, atten_weights, dim=1)
_, [h, c] = self.lstm_cell(inputs, [h, c])
preds = self.out(self.dropout(h))
x_t = self.embedding(preds.argmax(axis=1))
predictions.append(preds)
alphas.append(alpha)
predictions = nd.concat(*[x.expand_dims(axis=1) for x in predictions],
dim=1)
alphas = nd.concat(*[x.expand_dims(axis=1) for x in alphas], dim=1)
return predictions, alphas
def verifyLoadedModel(net, data):
data = nd.transpose(nd.array(data), (0, 3, 1, 2))
out = net(data)
predictions = nd.argmax(out, axis=1)
text_labels = [
't-shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt',
'sneaker', 'bag', 'ankle boot'
]
return [(int(p), text_labels[int(p)]) for p in predictions.asnumpy()]
def fn(heads, relations, tails, num_chunks, chunk_size, neg_sample_size):
hidden_dim = heads.shape[1]
emb_real, emb_img = nd.split(tails, num_outputs=2, axis=-1)
rel_real, rel_img = nd.split(relations, num_outputs=2, axis=-1)
real = emb_real * rel_real + emb_img * rel_img
img = -emb_real * rel_img + emb_img * rel_real
emb_complex = nd.concat(real, img, dim=-1)
tmp = emb_complex.reshape(num_chunks, chunk_size, hidden_dim)
heads = heads.reshape(num_chunks, neg_sample_size, hidden_dim)
heads = nd.transpose(heads, axes=(0, 2, 1))
return nd.linalg_gemm2(tmp, heads)
def postprocess(self, x):
"""
Description : module for postprocess
"""
output = F.relu(x)
output = self.conv_post_1(output)
output = F.relu(output)
output = self.conv_post_2(output)
output = nd.reshape(output, (output.shape[1], output.shape[2]))
output = F.transpose(output, axes=(1, 0))
return output
def hybrid_forward(self, F, x): # x(N, 1, C, 2)
# (n,1,c,2)->(n,k,c/k,2)
concat_step1 = x.reshape(
(x.shape[0], self._group, x.shape[2] / self._group, 2))
# (n,k,c/k,2)->(n,k,2,c/k)
concat_step2 = nd.transpose(concat_step1, (0, 1, 3, 2))
# (n,k,2,c/k)->(n, 1, 2*k,c/k)
concat_out = concat_step2.reshape(
(concat_step2.shape[0], 1, 2 * self._group,
concat_step2.shape[-1]))
return concat_out
def format_net_output(self, Y):
pred = nd.transpose(Y,(0,2,3,1)) #move channel to last dim
pred = pred.reshape((0,0,0,self.box_per_cell, self.numClass + 5)) # re-arrange last dim to two dim (B,())
#here you are responsible to define each field of output
predCls = nd.slice_axis(pred, begin=0, end=self.numClass,axis=-1)
predObj = nd.slice_axis(pred, begin=self.numClass, end=self.numClass+1,axis=-1)
predXY = nd.slice_axis(pred, begin=self.numClass+1, end=self.numClass+3,axis=-1)
predXY = nd.sigmoid(predXY)
predWH = nd.slice_axis(pred,begin=self.numClass+3,end=self.numClass+5,axis=-1)
XYWH = nd.concat(predXY, predWH, dim=-1)
return predCls, predObj, XYWH
def get_sample_point():
"""Grabs a single input/label pair from MNIST"""
def transform(data, label):
return data.astype(np.float32) / 255.0, label.astype(np.float32)
# Load ten random images from the test dataset
sample_data = mx.gluon.data.DataLoader(mx.gluon.data.vision.MNIST(
train=False, transform=transform),
1,
shuffle=True)
for data, label in sample_data:
img = nd.transpose(data, (1, 0, 2, 3))
img = nd.reshape(img, (28, 28, 1))
imtiles = nd.tile(img, (1, 1, 3))
plt.imshow(imtiles.asnumpy())
plt.savefig("test_input.png")
data = nd.transpose(data, (0, 3, 1, 2))
data = data.as_in_context(ctx).asnumpy()
label = int(label.asnumpy()[0])
return data, label
def transform(data, target_wd, target_ht):
# resize to target_wd * target_ht
data = mx.image.imresize(data, target_wd, target_ht)
# transpose from (target_wd, target_ht, 3)
# to (3, target_wd, target_ht)
data = nd.transpose(data, (2, 0, 1))
# normalize to [-1, 1]
data = data.astype(np.float32) / 127.5 - 1
# if image is greyscale, repeat 3 times to get RGB image.
if data.shape[0] == 1:
data = nd.tile(data, (3, 1, 1))
return data.reshape((1, ) + data.shape)
def predict(self, img):
img = nd.array(img)
img = nd.transpose(img, axes=(2, 0, 1)).astype('float32')
img = nd.expand_dims(img, axis=0)
#print(img.shape)
db = mx.io.DataBatch(data=(img, ))
self.model.forward(db, is_train=False)
net_out = self.model.get_outputs()
embedding = net_out[0].asnumpy()
embedding = sklearn.preprocessing.normalize(embedding)
return embedding
def reset_c2c(self):
self.select_triplets()
for identity,v in self.id2range.iteritems():
_list = range(*v)
for idx in _list:
s = imgrec.read_idx(idx)
ocontents.append(s)
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, args.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,))
self.mx_model.forward(db, is_train=False)
net_out = self.mx_model.get_outputs()
net_out = net_out[0].asnumpy()
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)
embedding = np.mean(embeddings, axis=0, keepdims=True)
embedding = sklearn.preprocessing.normalize(embedding)
sims = np.dot(embeddings, embedding).flatten()
assert len(sims)==len(_list)
for i in xrange(len(_list)):
_idx = _list[i]
self.idx2cos[_idx] = sims[i]
def get_embedding(args, imgrec, id, image_size, model):
s = imgrec.read_idx(id)
header, _ = mx.recordio.unpack(s)
ocontents = []
for idx in xrange(int(header.label[0]), int(header.label[1])):
#print('idx', idx)
s = imgrec.read_idx(idx)
ocontents.append(s)
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, args.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])
#print(header.label.shape, header.label)
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,))
db = mx.io.DataBatch(data=(data,))
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)
embedding = np.mean(embeddings, axis=0, keepdims=True)
embedding = sklearn.preprocessing.normalize(embedding).flatten()
return embedding
def load_bin(path, image_size):
bins, issame_list = pickle.load(open(path, 'rb'))
data_list = []
for flip in [0,1]:
data = nd.empty((len(issame_list)*2, 3, image_size[0], image_size[1]))
data_list.append(data)
for i in xrange(len(issame_list)*2):
_bin = bins[i]
img = mx.image.imdecode(_bin)
img = nd.transpose(img, axes=(2, 0, 1))
for flip in [0,1]:
if flip==1:
img = mx.ndarray.flip(data=img, axis=2)
data_list[flip][i][:] = img
if i%1000==0:
print('loading bin', i)
print(data_list[0].shape)
return (data_list, issame_list)
def load_dataset(lfw_dir, image_size):
lfw_pairs = read_pairs(os.path.join(lfw_dir, 'pairs.txt'))
lfw_paths, issame_list = get_paths(lfw_dir, lfw_pairs, 'jpg')
lfw_data_list = []
for flip in [0,1]:
lfw_data = nd.empty((len(lfw_paths), 3, image_size[0], image_size[1]))
lfw_data_list.append(lfw_data)
i = 0
for path in lfw_paths:
with open(path, 'rb') as fin:
_bin = fin.read()
img = mx.image.imdecode(_bin)
img = nd.transpose(img, axes=(2, 0, 1))
for flip in [0,1]:
if flip==1:
img = mx.ndarray.flip(data=img, axis=2)
lfw_data_list[flip][i][:] = img
i+=1
if i%1000==0:
print('loading lfw', i)
print(lfw_data_list[0].shape)
print(lfw_data_list[1].shape)
return (lfw_data_list, issame_list)
def parse_net_output(Y,numClass, box_per_cell):
pred = nd.transpose(Y,(0,2,3,1))
pred = pred.reshape((0,0,0,box_per_cell,numClass + 5)) #add one dim for boxes
predCls = nd.slice_axis(pred, begin = 0, end = numClass,axis=-1)
predObject = nd.slice_axis(pred,begin=numClass,end=numClass+1,axis=-1)
#predObject = nd.sigmoid(predObject)
predXY = nd.slice_axis(pred, begin = numClass + 1, end = numClass + 3, axis=-1)
#predXY = nd.sigmoid(predXY)
predWH = nd.slice_axis(pred, begin = numClass + 3, end = numClass + 5, axis=-1)
#x,y = convert_xy(predXY)
#w,h = convert_wh(predWH)
#w = nd.clip(w,0,1)
#h = nd.clip(h,0,1)
#x0 = nd.clip(x, 0, 1)
#y0 = nd.clip(y,0,1)
#x1 = nd.clip(x0 + w,0,1)
#y1 = np.clip(y0 + h, 0,1)
#x = x0
#y = y0
#w = x1 - x0
#h = y1 - y0
XYWH = nd.concat(predXY,predWH,dim=-1)
# pdb.set_trace()
return predCls, predObject, XYWH
def next(self):
"""Returns the next batch of data."""
#print('next')
batch_size = self.batch_size
batch_data = nd.empty((batch_size,)+self.data_shape)
batch_label = nd.empty((batch_size,)+self.label_shape)
i = 0
#self.cutoff = random.randint(800,1280)
try:
while i < batch_size:
#print('N', i)
data, label = self.next_sample()
data = nd.array(data)
data = nd.transpose(data, axes=(2, 0, 1))
label = nd.array(label)
#print(data.shape, label.shape)
batch_data[i][:] = data
batch_label[i][:] = label
i += 1
except StopIteration:
if not i:
raise StopIteration
return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)
def postprocess_data(self, datum):
"""Final postprocessing step before image is loaded into the batch."""
return nd.transpose(datum, axes=(2, 0, 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
end_epoch = args.end_epoch
pretrained = '../model/resnet-152'
load_epoch = args.load_epoch
args.image_size = 160
per_batch_size = 60
args.ctx_num = len(ctx)
args.batch_size = per_batch_size*args.ctx_num
#args.all_batch_size = args.batch_size*args.ctx_num
args.bag_size = 3600
args.margin = 0.2
args.num_classes = 10575 #webface
data_shape = (3,args.image_size,args.image_size)
begin_epoch = 0
base_lr = 0.05
base_wd = 0.0002
base_mom = 0.0
lr_decay = 0.98
if not args.retrain:
#load and initialize params
print(pretrained)
_, arg_params, aux_params = mx.model.load_checkpoint(pretrained, load_epoch)
sym, arg_params, aux_params = get_symbol(args, arg_params, aux_params)
#arg_params, aux_params = load_param(pretrained, epoch, convert=True)
data_shape_dict = {'data': (args.batch_size, 3, args.image_size, args.image_size), 'softmax_label': (args.batch_size,)}
resnet_dcn.init_weights(sym, data_shape_dict, arg_params, aux_params)
else:
pretrained = args.prefix
sym, arg_params, aux_params = mx.model.load_checkpoint(pretrained, load_epoch)
begin_epoch = load_epoch
end_epoch = begin_epoch+10
base_wd = 0.00005
lr_decay = 0.5
base_lr = 0.015
# infer max shape
model = mx.mod.Module(
context = ctx,
symbol = sym,
#label_names = [],
#fixed_param_prefix = fixed_param_prefix,
)
train_dataiter = FaceIter(
path_imglist = "/raid5data/dplearn/faceinsight_align_webface.lst",
data_shape = data_shape,
mod = model,
ctx_num = args.ctx_num,
batch_size = args.batch_size,
bag_size = args.bag_size,
images_per_person = 5,
)
#_dice = DiceMetric()
_acc = AccMetric()
eval_metrics = [mx.metric.create(_acc)]
# rpn_eval_metric, rpn_cls_metric, rpn_bbox_metric, eval_metric, cls_metric, bbox_metric
#for child_metric in [fcn_loss_metric]:
# eval_metrics.add(child_metric)
# callback
#batch_end_callback = callback.Speedometer(input_batch_size, frequent=args.frequent)
#epoch_end_callback = mx.callback.module_checkpoint(mod, prefix, period=1, save_optimizer_states=True)
# decide learning rate
#lr_step = '10,20,30'
#train_size = 4848
#nrof_batch_in_epoch = int(train_size/input_batch_size)
#print('nrof_batch_in_epoch:', nrof_batch_in_epoch)
#lr_factor = 0.1
#lr_epoch = [float(epoch) for epoch in lr_step.split(',')]
#lr_epoch_diff = [epoch - begin_epoch for epoch in lr_epoch if epoch > begin_epoch]
#lr = base_lr * (lr_factor ** (len(lr_epoch) - len(lr_epoch_diff)))
#lr_iters = [int(epoch * train_size / batch_size) for epoch in lr_epoch_diff]
#print 'lr', lr, 'lr_epoch_diff', lr_epoch_diff, 'lr_iters', lr_iters
#lr_scheduler = MultiFactorScheduler(lr_iters, lr_factor)
# optimizer
#optimizer_params = {'momentum': 0.9,
# 'wd': 0.0005,
# 'learning_rate': base_lr,
# 'rescale_grad': 1.0,
# 'clip_gradient': None}
initializer = mx.init.Xavier(rnd_type='gaussian', factor_type="in", magnitude=2)
#opt = optimizer.SGD(learning_rate=base_lr, momentum=0.9, wd=base_wd, rescale_grad=(1.0/args.batch_size))
opt = optimizer.SGD(learning_rate=base_lr, momentum=base_mom, wd=base_wd, rescale_grad=1.0)
#opt = optimizer.AdaGrad(learning_rate=base_lr, wd=base_wd, rescale_grad=1.0)
_cb = mx.callback.Speedometer(args.batch_size, 10)
lfw_dir = '/raid5data/dplearn/lfw_mtcnn'
lfw_pairs = lfw.read_pairs(os.path.join(lfw_dir, 'pairs.txt'))
lfw_paths, issame_list = lfw.get_paths(lfw_dir, lfw_pairs, 'png')
imgs = []
lfw_data_list = []
for flip in [0,1]:
lfw_data = nd.empty((len(lfw_paths), 3, args.image_size, args.image_size))
i = 0
for path in lfw_paths:
with open(path, 'rb') as fin:
_bin = fin.read()
img = mx.image.imdecode(_bin)
img = nd.transpose(img, axes=(2, 0, 1))
if flip==1:
img = img.asnumpy()
for c in xrange(img.shape[0]):
img[c,:,:] = np.fliplr(img[c,:,:])
img = nd.array( img )
#print(img.shape)
lfw_data[i][:] = img
i+=1
if i%1000==0:
print('loading lfw', i)
print(lfw_data.shape)
lfw_data_list.append(lfw_data)
def lfw_test(nbatch):
print('testing lfw..')
embeddings_list = []
for i in xrange( len(lfw_data_list) ):
lfw_data = lfw_data_list[i]
embeddings = None
ba = 0
while ba<lfw_data.shape[0]:
bb = min(ba+args.batch_size, lfw_data.shape[0])
_data = nd.slice_axis(lfw_data, axis=0, begin=ba, end=bb)
_label = nd.ones( (bb-ba,) )
db = mx.io.DataBatch(data=(_data,), label=(_label,))
model.forward(db, is_train=False)
net_out = model.get_outputs()
_embeddings = net_out[0].asnumpy()
if embeddings is None:
embeddings = np.zeros( (lfw_data.shape[0], _embeddings.shape[1]) )
embeddings[ba:bb,:] = _embeddings
ba = bb
embeddings_list.append(embeddings)
acc_list = []
embeddings = embeddings_list[0]
_, _, accuracy, val, val_std, far = lfw.evaluate(embeddings, issame_list, nrof_folds=10)
acc_list.append(np.mean(accuracy))
print('[%d]Accuracy: %1.3f+-%1.3f' % (nbatch, np.mean(accuracy), np.std(accuracy)))
print('Validation rate: %2.5f+-%2.5f @ FAR=%2.5f' % (val, val_std, far))
embeddings = np.concatenate(embeddings_list, axis=1)
embeddings = sklearn.preprocessing.normalize(embeddings)
print(embeddings.shape)
_, _, accuracy, val, val_std, far = lfw.evaluate(embeddings, issame_list, nrof_folds=10)
acc_list.append(np.mean(accuracy))
print('[%d]Accuracy-Flip: %1.3f+-%1.3f' % (nbatch, np.mean(accuracy), np.std(accuracy)))
print('Validation rate: %2.5f+-%2.5f @ FAR=%2.5f' % (val, val_std, far))
pca = PCA(n_components=128)
embeddings = pca.fit_transform(embeddings)
embeddings = sklearn.preprocessing.normalize(embeddings)
print(embeddings.shape)
_, _, accuracy, val, val_std, far = lfw.evaluate(embeddings, issame_list, nrof_folds=10)
acc_list.append(np.mean(accuracy))
print('[%d]Accuracy-PCA: %1.3f+-%1.3f' % (nbatch, np.mean(accuracy), np.std(accuracy)))
print('Validation rate: %2.5f+-%2.5f @ FAR=%2.5f' % (val, val_std, far))
return max(*acc_list)
#global_step = 0
highest_acc = [0.0]
last_save_acc = [0.0]
def _batch_callback(param):
#global global_step
mbatch = param.nbatch+1
if mbatch % 4000 == 0:
opt.lr *= lr_decay
#print(param.nbatch, opt.lr)
_cb(param)
if param.nbatch%100==0:
print('lr-batch-epoch:',opt.lr,param.nbatch,param.epoch)
if param.nbatch>=0 and param.nbatch%400==0:
acc = lfw_test(param.nbatch)
if acc>highest_acc[0]:
highest_acc[0] = acc
if acc>0.9 and acc-last_save_acc[0]>=0.01:
print('saving', mbatch, acc, last_save_acc[0])
_arg, _aux = model.get_params()
mx.model.save_checkpoint(args.prefix, mbatch, model.symbol, _arg, _aux)
last_save_acc[0] = acc
print('[%d]highest Accu: %1.3f'%(param.nbatch, highest_acc[0]))
sys.stdout.flush()
sys.stderr.flush()
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 )
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 transform_mnist(data, label):
return nd.transpose(data.astype('float32'), (2,0,1)/255, label.astype('float32'));
def next(self):
"""Returns the next batch of data."""
#print('next')
batch_size = self.batch_size
batch_data = nd.empty((batch_size,)+self.data_shape)
batch_label = nd.empty((batch_size,)+self.label_shape)
if self.use_coherent:
batch_label2 = nd.empty((batch_size,)+self.label_shape)
batch_coherent_label = nd.empty((batch_size,6))
i = 0
#self.cutoff = random.randint(800,1280)
try:
while i < batch_size:
#print('N', i)
data, label, annot = self.next_sample()
if not self.use_coherent:
R = self.do_aug(data, label, annot)
if R is None:
continue
data, label = R
#data, label, data2, label2, M = R
#ind = np.unravel_index(np.argmax(label[0], axis=None), label[0].shape)
#print(label.shape, np.count_nonzero(label[0]), ind)
#print(label[0,25:35,0:10])
data = nd.array(data)
data = nd.transpose(data, axes=(2, 0, 1))
label = nd.array(label)
#print(data.shape, label.shape)
try:
self.check_valid_image(data)
except RuntimeError as e:
logging.debug('Invalid image, skipping: %s', str(e))
continue
batch_data[i][:] = data
batch_label[i][:] = label
i += 1
else:
R = self.do_aug(data, label, annot)
if R is None:
continue
data, label, data2, label2, M = R
data = nd.array(data)
data = nd.transpose(data, axes=(2, 0, 1))
label = nd.array(label)
data2 = nd.array(data2)
data2 = nd.transpose(data2, axes=(2, 0, 1))
label2 = nd.array(label2)
M = nd.array(M)
#print(data.shape, label.shape)
try:
self.check_valid_image(data)
except RuntimeError as e:
logging.debug('Invalid image, skipping: %s', str(e))
continue
batch_data[i][:] = data
batch_label[i][:] = label
#batch_label2[i][:] = label2
batch_coherent_label[i][:] = M
#i+=1
j = i+self.per_batch_size//2
batch_data[j][:] = data2
batch_label[j][:] = label2
batch_coherent_label[j][:] = M
i += 1
if j%self.per_batch_size==self.per_batch_size-1:
i = j+1
except StopIteration:
if not i:
raise StopIteration
#return {self.data_name : batch_data,
# self.label_name : batch_label}
#print(batch_data.shape, batch_label.shape)
if not self.use_coherent:
return mx.io.DataBatch([batch_data], [batch_label], batch_size - i)
else:
#return mx.io.DataBatch([batch_data], [batch_label, batch_label2, batch_coherent_label], batch_size - i)
return mx.io.DataBatch([batch_data], [batch_label, batch_coherent_label], batch_size - i)
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 transform_mnist(data, label):
# change data from height x weight x channel to channel x height x weight
return nd.transpose(data.astype('float32'), (2,0,1))/255, label.astype('float32')
