train.py

import os
import sys
import math
import time
import argparse
import functools
import numpy as np
import paddle
import paddle.fluid as fluid
import models
from losses import tripletloss
from losses import quadrupletloss
from losses import emlloss
from losses.metrics import recall_topk
from utility import add_arguments, print_arguments

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('train_batch_size', int, 80, "Minibatch size.")
add_arg('test_batch_size', int, 10, "Minibatch size.")
add_arg('num_epochs', int, 120, "number of epochs.")
add_arg('image_shape', str, "3,224,224", "input image size")
add_arg('model_save_dir', str, "output", "model save directory")
add_arg('with_mem_opt', bool, True,
        "Whether to use memory optimization or not.")
add_arg('pretrained_model', str, None, "Whether to use pretrained model.")
add_arg('checkpoint', str, None, "Whether to resume checkpoint.")
add_arg('lr', float, 0.1, "set learning rate.")
add_arg('lr_strategy', str, "piecewise_decay",
        "Set the learning rate decay strategy.")
add_arg('model', str, "SE_ResNeXt50_32x4d", "Set the network to use.")
add_arg('loss_name', str, "tripletloss", "Set the loss type to use.")
add_arg('samples_each_class', int, 2, "Samples each class.")
add_arg('margin', float, 0.1, "margin.")
add_arg('alpha', float, 0.0, "alpha.")
# yapf: enable

model_list = [m for m in dir(models) if "__" not in m]

def optimizer_setting(params):
    ls = params["learning_strategy"]
    assert ls["name"] == "piecewise_decay", \
           "learning rate strategy must be {}, \
           but got {}".format("piecewise_decay", lr["name"])

    step = 10000
    bd = [step * e for e in ls["epochs"]]
    base_lr = params["lr"]
    lr = []
    lr = [base_lr * (0.1 ** i) for i in range(len(bd) + 1)]
    optimizer = fluid.optimizer.Momentum(
        learning_rate=fluid.layers.piecewise_decay(
            boundaries=bd, values=lr),
        momentum=0.9,
        regularization=fluid.regularizer.L2Decay(1e-4))

    return optimizer

def train(args):
    # parameters from arguments
    model_name = args.model
    checkpoint = args.checkpoint
    pretrained_model = args.pretrained_model
    with_memory_optimization = args.with_mem_opt
    model_save_dir = args.model_save_dir
    loss_name = args.loss_name

    image_shape = [int(m) for m in args.image_shape.split(",")]

    assert model_name in model_list, "{} is not in lists: {}".format(args.model, model_list)

    image = fluid.layers.data(name='image', shape=image_shape, dtype='float32')
    label = fluid.layers.data(name='label', shape=[1], dtype='int64')
    
    # model definition
    model = models.__dict__[model_name]()
    out = model.net(input=image, class_dim=200)

    if loss_name == "tripletloss":
        metricloss = tripletloss(
                train_batch_size = args.train_batch_size, 
                margin=args.margin)
        cost_metric = metricloss.loss(out[0])
        avg_cost_metric = fluid.layers.mean(x=cost_metric)
    elif loss_name == "quadrupletloss":
        metricloss = quadrupletloss(
                train_batch_size = args.train_batch_size,
                samples_each_class = args.samples_each_class,
                margin=args.margin)
        cost_metric = metricloss.loss(out[0])
        avg_cost_metric = fluid.layers.mean(x=cost_metric)
    elif loss_name == "emlloss":
        metricloss = emlloss(
                train_batch_size = args.train_batch_size, 
                samples_each_class = args.samples_each_class
        )
        cost_metric = metricloss.loss(out[0])
        avg_cost_metric = fluid.layers.mean(x=cost_metric)

    cost_cls = fluid.layers.cross_entropy(input=out[1], label=label)
    avg_cost_cls = fluid.layers.mean(x=cost_cls)
    acc_top1 = fluid.layers.accuracy(input=out[1], label=label, k=1)
    acc_top5 = fluid.layers.accuracy(input=out[1], label=label, k=5)
    avg_cost = avg_cost_metric + args.alpha*avg_cost_cls
    

    test_program = fluid.default_main_program().clone(for_test=True)

    # parameters from model and arguments
    params = model.params    
    params["lr"] = args.lr
    params["num_epochs"] = args.num_epochs
    params["learning_strategy"]["batch_size"] = args.train_batch_size
    params["learning_strategy"]["name"] = args.lr_strategy

    # initialize optimizer
    optimizer = optimizer_setting(params)
    opts = optimizer.minimize(avg_cost)

    global_lr = optimizer._global_learning_rate()

    place = fluid.CUDAPlace(0)
    exe = fluid.Executor(place)
    exe.run(fluid.default_startup_program())

    if checkpoint is not None:
        fluid.io.load_persistables(exe, checkpoint)

    if pretrained_model:
        assert(checkpoint is None)
        def if_exist(var):
            has_var = os.path.exists(os.path.join(pretrained_model, var.name))
            if has_var:
                print('var: %s found' % (var.name))
            return has_var
        fluid.io.load_vars(exe, pretrained_model, predicate=if_exist)

    train_reader = paddle.batch(metricloss.train_reader, batch_size=args.train_batch_size)
    test_reader = paddle.batch(metricloss.test_reader, batch_size=args.test_batch_size)
    feeder = fluid.DataFeeder(place=place, feed_list=[image, label])

    train_exe = fluid.ParallelExecutor(use_cuda=True, loss_name=avg_cost.name)

    fetch_list_train = [avg_cost_metric.name, avg_cost_cls.name, acc_top1.name, acc_top5.name, global_lr.name]
    fetch_list_test = [out[0].name]

    if with_memory_optimization:
        fluid.memory_optimize(fluid.default_main_program(), skip_opt_set=set(fetch_list_train))

    for pass_id in range(params["num_epochs"]):
        train_info = [[], [], [], []]
        for batch_id, data in enumerate(train_reader()):
            t1 = time.time()
            loss_metric, loss_cls, acc1, acc5, lr = train_exe.run(fetch_list_train, feed=feeder.feed(data))
            t2 = time.time()
            period = t2 - t1
            loss_metric = np.mean(np.array(loss_metric))
            loss_cls = np.mean(np.array(loss_cls))
            acc1 = np.mean(np.array(acc1))
            acc5 = np.mean(np.array(acc5))
            lr = np.mean(np.array(lr))
            train_info[0].append(loss_metric)
            train_info[1].append(loss_cls)
            train_info[2].append(acc1)
            train_info[3].append(acc5)
            if batch_id % 10 == 0:
                print("Pass {0}, trainbatch {1}, lr {2}, loss_metric {3}, loss_cls {4}, acc1 {5}, acc5 {6}, time {7}".format(pass_id,  \
                      batch_id, lr, loss_metric, loss_cls, acc1, acc5, "%2.2f sec" % period))

        train_loss_metric = np.array(train_info[0]).mean()
        train_loss_cls = np.array(train_info[1]).mean()
        train_acc1 = np.array(train_info[2]).mean()
        train_acc5 = np.array(train_info[3]).mean()
        f = []
        l = []
        for batch_id, data in enumerate(test_reader()):
            if len(data) < args.test_batch_size:
                continue
            t1 = time.time()
            [feas] = exe.run(test_program, fetch_list = fetch_list_test, feed=feeder.feed(data))
            label = np.asarray([x[1] for x in data])
            f.append(feas)
            l.append(label)

            t2 = time.time()
            period = t2 - t1
            if batch_id % 20 == 0:
                print("Pass {0}, testbatch {1}, time {2}".format(pass_id,  \
                      batch_id, "%2.2f sec" % period))

        f = np.vstack(f)
        l = np.hstack(l)
        recall = recall_topk(f, l, k = 1)
        print("End pass {0}, train_loss_metric {1}, train_loss_cls {2}, train_acc1 {3}, train_acc5 {4}, test_recall {5}".format(pass_id,  \
              train_loss_metric, train_loss_cls, train_acc1, train_acc5, recall))
        sys.stdout.flush()

        model_path = os.path.join(model_save_dir + '/' + model_name,
                                  str(pass_id))
        if not os.path.isdir(model_path):
            os.makedirs(model_path)
        fluid.io.save_persistables(exe, model_path)

def main():
    args = parser.parse_args()
    print_arguments(args)
    train(args)

if __name__ == '__main__':
    main()