insightface中recognition训练过程

insightface中recognition训练过程

1 config文件

到目录/recognition/下有sample_config.py文件，这是给了个样例配置文件，我们复制一下起名config.py然后编辑 config.py

cp sample_config.py config.py
vim config.py

config下做了一些配置

default默认配置
config配置
network网络模型参数
loss损失函数选择
dataset数据集配置
generate_config生成配置文件，把network、loss、dataset配置都加载config上。

然后看train.py文件

2 main

主函数比较简单，参数填进去就可以了。这里需要特别注意下参数这一块。

def main():
    global args
    args = parse_args()
    train_net(args)

到最开始那里有参数，如下，我们注意先来三个参数dataset、network、loss然后有一个generate_config生成配置，吧三个参数组合到config里，这里这三个参数要和config.py里面每个参数起的名字一样哦
例如dataset : emore retina 当然你也可以自己改配置文件自己加
network :r100 r50 r100fc
…

def parse_args():
  parser = argparse.ArgumentParser(description='Train face network')
  # general
  parser.add_argument('--dataset', default=default.dataset, help='dataset config')
  parser.add_argument('--network', default=default.network, help='network config')
  parser.add_argument('--loss', default=default.loss, help='loss config')
  args, rest = parser.parse_known_args()
  generate_config(args.network, args.dataset, args.loss)
  parser.add_argument('--models-root', default=default.models_root, help='root directory to save model.')
  parser.add_argument('--pretrained', default=default.pretrained, help='pretrained model to load')
  parser.add_argument('--pretrained-epoch', type=int, default=default.pretrained_epoch, help='pretrained epoch to load')
  parser.add_argument('--ckpt', type=int, default=default.ckpt, help='checkpoint saving option. 0: discard saving. 1: save when necessary. 2: always save')
  parser.add_argument('--verbose', type=int, default=default.verbose, help='do verification testing and model saving every verbose batches')
  parser.add_argument('--lr', type=float, default=default.lr, help='start learning rate')
  parser.add_argument('--lr-steps', type=str, default=default.lr_steps, help='steps of lr changing')
  parser.add_argument('--wd', type=float, default=default.wd, help='weight decay')
  parser.add_argument('--mom', type=float, default=default.mom, help='momentum')
  parser.add_argument('--frequent', type=int, default=default.frequent, help='')
  parser.add_argument('--per-batch-size', type=int, default=default.per_batch_size, help='batch size in each context')
  parser.add_argument('--kvstore', type=str, default=default.kvstore, help='kvstore setting')
  args = parser.parse_args()
  return args

然后主要看train.py文件

3 train

前面会选择GPU，然后make 存模型的文件夹，然后读数据集bin文件

加载预训练模型，如果有预训练模型，就直接加载checkpoint模型参数到arg_params，aux_params，然后sym = get_symbol(args)走一个前向传播，如果没有模型则声明个none。参数在后面fit时有用。

    if len(args.pretrained)==0:
      arg_params = None
      aux_params = None
      sym = get_symbol(args)
      if config.net_name=='spherenet':
        data_shape_dict = {'data' : (args.per_batch_size,)+data_shape}
        spherenet.init_weights(sym, data_shape_dict, args.num_layers)
    else:
      print('loading', args.pretrained, args.pretrained_epoch)
      _, arg_params, aux_params = mx.model.load_checkpoint(args.pretrained, args.pretrained_epoch)
      sym = get_symbol(args)

这里有个get_symbol前面有定义，我们到前面能看到，写了一大堆，我们看一点
这里先获取你的网络的模型net_name这一套在/recognition/symbol有resnet等等的get_symbol，这个特征提完后，加全连接层，分别有softmax等等。

def get_symbol(args):
  embedding = eval(config.net_name).get_symbol()
  all_label = mx.symbol.Variable('softmax_label')
  gt_label = all_label
  is_softmax = True
  if config.loss_name=='softmax': #softmax
    _weight = mx.symbol.Variable("fc7_weight", shape=(config.num_classes, config.emb_size), 
        lr_mult=config.fc7_lr_mult, wd_mult=config.fc7_wd_mult, init=mx.init.Normal(0.01))
    if config.fc7_no_bias:
      fc7 = mx.sym.FullyConnected(data=embedding, weight = _weight, no_bias = True, num_hidden=config.num_classes, name='fc7')
    else:
      _bias = mx.symbol.Variable('fc7_bias', lr_mult=2.0, wd_mult=0.0)
      fc7 = mx.sym.FullyConnected(data=embedding, weight = _weight, bias = _bias, num_hidden=config.num_classes, name='fc7')

回到train上来
然后计算算力，这部分时flops_counter.py在/common下，有兴趣可以看一下。

    #计算消耗算力
    if config.count_flops:
      all_layers = sym.get_internals()
      _sym = all_layers['fc1_output']
      FLOPs = flops_counter.count_flops(_sym, data=(1,3,image_size[0],image_size[1]))
      _str = flops_counter.flops_str(FLOPs)
      print('Network FLOPs: %s'%_str)

后面创建模型
contex默认选择CPU，给ctxGPU序号
symbol给出计算的输出即可。
symbol: the network definition
context: the device (or a list of devices) to use for execution
data_names : the list of input data variable names
label_names : the list of input label variable names

    model = mx.mod.Module(
        context       = ctx,
        symbol        = sym,
    )

主要看model.fit

    model.fit(train_dataiter,
        begin_epoch        = begin_epoch,
        num_epoch          = 999999,
        eval_data          = val_dataiter,
        eval_metric        = eval_metrics,
        kvstore            = args.kvstore,
        optimizer          = opt,
        #optimizer_params   = optimizer_params,
        initializer        = initializer,
        arg_params         = arg_params,
        aux_params         = aux_params,
        allow_missing      = True,
        batch_end_callback = _batch_callback,
        epoch_end_callback = epoch_cb )

参数解释如下
train_data (DataIter) – Train DataIter.
eval_data (DataIter) – 验证集数据，每个epoch会验证一下，可以选择none
epoch_end_callback (function or list of functions) – 每一次epoch调用一次
batch_end_callback (function or list of function) – 每个batch调用一次，可以存放一些打印函数或者我们更新学习率
kvstore (str or KVStore) – Defaults to ‘local’.
optimizer (str or Optimizer) – Defaults to ‘sgd’.优化器
initializer (Initializer) – The initializer is called to initialize the module parameters when they are not already initialized.
arg_params (dict) – 初始化参数，默认没有
aux_params (dict) – aux参数同上
allow_missing (bool) – 在有预训练参数时是否允许丢失，也就是如果允许，则上面的参数万一有丢的，我们就随机初始化。
begin_epoch (int) – 开始的次数，从几开始计数
num_epoch (int) – 训练epoch

然后看到model每个epoch调用了回调函数，做了点啥呢
主要时当数据量达到一定就打印数据，当测试效果好的时候就保存模型。

4运行

CUDA_VISIBLE_DEVICES='0,1,2,3' python -u train.py --network r50 --loss arcface --dataset emore