def test():
"""
Test model accuracy on test dataset.
测试模型在测试集上的准确率。
"""
print("Start to test...")
ctx = mx.gpu() if args.use_gpu else mx.cpu()
_, test_iter = data_loader(args.batch_size)
model = LeNetPlus()
model.load_parameters(os.path.join(args.ckpt_dir,
args.prefix + "-best.params"),
ctx=ctx,
allow_missing=True)
#
center_net = CenterLoss(num_classes=args.num_classes,
feature_size=args.feature_size,
lmbd=args.lmbd,
ctx=mx.cpu())
center_net.load_parameters(os.path.join(
args.ckpt_dir, args.prefix + "-feature_matrix.params"),
ctx=ctx)
start_time = time.time()
test_accuracy, features, predictions, labels = evaluate_accuracy(
test_iter, model, center_net, args.eval_method, ctx)
elapsed_time = time.time() - start_time
print("Test_acc: %s, Elapsed_time: %f s" % (test_accuracy, elapsed_time))
# make directory
if not os.path.exists(args.out_dir):
os.makedirs(args.out_dir)
# draw feature map
if args.plotting:
plot_features(features,
labels,
num_classes=args.num_classes,
fpath=os.path.join(args.out_dir, "%s.png" % args.prefix))
def second_train():
"""
Train a model using only center loss based on pretrained model.
In order to avoid feature matrix becoming zero matrix, fix CenterLoss' parameters not to train it.
基于之前训练的模型，仅使用center loss对模型进行训练。
为了避免在训练的过程，CenterLoss中的特征矩阵变为0矩阵，将其参数固定不对其进行训练。
"""
print("Start to train LeNet++ with CenterLoss...")
train_iter, test_iter = data_loader(args.batch_size)
ctx = mx.gpu() if args.use_gpu else mx.cpu()
# main model (LeNetPlus), loss, trainer
model = LeNetPlus(classes=args.num_classes, feature_size=args.feature_size)
model.load_parameters(os.path.join(args.ckpt_dir,
args.prefix + "-best.params"),
ctx=ctx,
allow_missing=True)
network_trainer = gluon.Trainer(model.collect_params(),
optimizer="sgd",
optimizer_params={
"learning_rate": args.lr,
"wd": args.wd
})
center_loss = CenterLoss(num_classes=args.num_classes,
feature_size=args.feature_size,
lmbd=args.lmbd,
ctx=ctx)
center_loss.load_parameters(os.path.join(
args.ckpt_dir, args.prefix + "-feature_matrix.params"),
ctx=ctx)
center_loss.params.setattr("grad_req", "null") #
smoothing_constant, moving_loss = .01, 0.0
best_acc = 0.0
for epoch in range(args.epochs):
# using learning rate decay during training process
if (epoch > 0) and (epoch % args.lr_step == 0):
network_trainer.set_learning_rate(network_trainer.learning_rate *
args.lr_factor)
start_time = time.time()
for i, (data, label) in enumerate(train_iter):
data = data.as_in_context(ctx)
label = label.as_in_context(ctx)
# only calculate loss with center loss
with autograd.record():
output, features = model(data)
loss = center_loss(features, label)
loss.backward()
# only update parameters of LeNet++
network_trainer.step(
args.batch_size, ignore_stale_grad=True
) # without ignore_stale_grad=True, it will raise warning information
# 去除ignore_stale_grad=True后，模型训练会出错，梯度无法进行更新
# calculate smoothed loss value
curr_loss = nd.mean(loss).asscalar()
moving_loss = (curr_loss if ((i == 0) and (epoch == 0)) else
(1 - smoothing_constant) * moving_loss +
smoothing_constant * curr_loss)
# training cost time
elapsed_time = time.time() - start_time
train_accuracy, train_ft, _, train_lb = evaluate_accuracy(
train_iter, model, center_loss, args.eval_method, ctx)
test_accuracy, test_ft, _, test_lb = evaluate_accuracy(
test_iter, model, center_loss, args.eval_method, ctx)
# draw feature map with different prefix, to make it convenient to compare features
if args.plotting:
plot_features(train_ft,
train_lb,
num_classes=args.num_classes,
fpath=os.path.join(
args.out_dir, "%s-second-train-epoch-%d.png" %
(args.prefix, epoch)))
plot_features(test_ft,
test_lb,
num_classes=args.num_classes,
fpath=os.path.join(
args.out_dir, "%s-second-test-epoch-%d.png" %
(args.prefix, epoch)))
logging.warning(
"Epoch [%d]: Loss=%f, Train-Acc=%f, Test-Acc=%f, Epoch-time=%f" %
(epoch, moving_loss, train_accuracy, test_accuracy, elapsed_time))
if test_accuracy > best_acc:
best_acc = test_accuracy
model.save_parameters(
os.path.join(args.ckpt_dir,
args.prefix + "-second-best.params"))
