深度学习【22】Mxnet多任务（multi-task）训练

github上有两个版本的多任务训练分别是：
1、https://github.com/miraclewkf/multi-task-MXNet
2、mxnet自带的例子
第一个由于其数据迭代器是Image，可能会比较慢。
第二个的例子是mnist，需要自己修改数据迭代器。
这里主要记录基于ImageRecordIter迭代器的多任务训练。

1、数据制作
需要自己生成*.lst文件，里面内容如下：

index   task1标签   task2标签    task3标签    图片路径(这行是说明，不需要写入，每一列用\t隔开)
2476    0.000000    0.000000    1.000000    photo_02_8159/00022552.jpg 
7623    3.000000    2.000000    2.000000    photo_03_7397/00029434.jpg
14149   0.000000    0.000000    1.000000    photo_05_15560/00060839.jpg
6874    3.000000    1.000000    2.000000    photo_03_7397/00028414.jpg
6048    0.000000    0.000000    1.000000    photo_02_8159/00027259.jpg
14479   3.000000    3.000000    2.000000    photo_05_15560/00065068.jpg
10429   2.000000    0.000000    1.000000    photo_04_15224/00040186.jpg
6949    3.000000    0.000000    1.000000    photo_03_7397/00028521.jpg
81      3.000000    3.000000    2.000000    photo_01_19992/00002536.jpg
11725   2.000000    0.000000    1.000000    photo_05_15560/00051778.jpg
1517    2.000000    3.000000    2.000000    photo_02_8159/00021245.jpg

具体是生成方法可以参考mxnet提供的im2rec.py，可以自己写一个make_list函数。
生成*.rec文件。这个文件可以用im2rec.py生成，同时需要把pack-label设置为True。
2、修改模型结构
添加3个mx.symbol.SoftmaxOutput损失函数（因为我这边是3个任务）：

    fc1 = mx.symbol.FullyConnected(data=flat, num_hidden=5, name='fc1') #任务1 有5个类别
    fc2 = mx.symbol.FullyConnected(data=flat, num_hidden=15, name='fc2') #任务2 有15个类别
    fc3 = mx.symbol.FullyConnected(data=flat, num_hidden=3, name='fc3') #任务3 有3个类别
    #分别为这三个任务添加softmax损失函数，注意每个函数的名称，后面会用到
    s1 = mx.symbol.SoftmaxOutput(data=fc1, name='softmax1') 
    s2 = mx.symbol.SoftmaxOutput(data=fc2, name='softmax2')
    s3 = mx.symbol.SoftmaxOutput(data=fc3, name='softmax3')
    return  mx.symbol.Group([s1,s2,s3])

3、编写ImageRecordIter选项

    train = mx.io.ImageRecordIter(
        path_imgrec='/path/to/train.rec',
        label_name=['softmax1_label', 'softmax2_label', 'softmax3_label'],#label名称，于softmax名称一样，后面要加入_label
        label_width=3, #重要，需要设置label宽度为3，因为有3个任务
        data_shape=[3,224,224],
        batch_size=64
    )

    val = mx.io.ImageRecordIter(
        path_imgrec='/path/to/val.rec',
        label_name=['softmax1_label', 'softmax2_label', 'softmax3_label'],
        label_width=3,
        batch_size=64,
        data_shape=[3,224,224],
    )

4、定义多任务训练迭代器

class MultiTask_iter(mx.io.DataIter):
    def __init__(self, data_iter):
        super(MultiTask_iter,self).__init__('multitask_iter')
        self.data_iter = data_iter
        self.batch_size = self.data_iter.batch_size

    @property
    def provide_data(self):
        return self.data_iter.provide_data

    @property
    def provide_label(self):
        provide_label = self.data_iter.provide_label[0]
        # the name of the label if corresponding to the model you define in get_fine_tune_model() function
        return [('softmax1_label', [provide_label[1][0]]),#需要注意的地方
        ('softmax2_label', [provide_label[1][0]]),
        ('softmax3_label', [provide_label[1][0]])]

    def hard_reset(self):
        self.data_iter.hard_reset()

    def reset(self):
        self.data_iter.reset()

    def next(self):
        batch = self.data_iter.next()
        #需要注意的地方
        label = batch.label[0]
        ll = label.asnumpy()
        label1 = mx.nd.array(ll[:,0]).astype('float32')
        label2 = mx.nd.array(ll[:,1]).astype('float32')
        label3 = mx.nd.array(ll[:,2]).astype('float32')
        # we set task 2 as: if label>0 or not

        return mx.io.DataBatch(data=batch.data, label=[label1,label2,label3], \
                pad=batch.pad, index=batch.index)

5、定义正确率计算方法

class Multi_Accuracy(mx.metric.EvalMetric):
    """Calculate accuracies of multi label"""

    def __init__(self, num=None):
        super(Multi_Accuracy, self).__init__('multi-accuracy', num)

    def update(self, labels, preds):
        mx.metric.check_label_shapes(labels, preds)

        if self.num is not None:
            assert len(labels) == self.num

        for i in range(len(labels)):
            pred_label = mx.nd.argmax_channel(preds[i]).asnumpy().astype('int32')
            label = labels[i].asnumpy().astype('int32')

            mx.metric.check_label_shapes(label, pred_label)

            if i is None:
                self.sum_metric += (pred_label.flat == label.flat).sum()
                self.num_inst += len(pred_label.flat)
            else:
                self.sum_metric[i] += (pred_label.flat == label.flat).sum()
                self.num_inst[i] += len(pred_label.flat)

6、训练

    train = MultiTask_iter(train)#调用多任务迭代器，其中train参数就是第3步的东西
    val = MultiTask_iter(val)

    new_sym = get_symbol(10,50,image_shape)

    optimizer_params = {
            'learning_rate': 0.001,
            'momentum' : args.mom,
            'wd' : args.wd,
           }
    initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)
    model = mx.mod.Module(
        context       = devs,
        symbol        = new_sym,
        data_names=['data'],
        label_names=['softmax1_label','softmax2_label','softmax3_label']
    )
    saveroot = args.save_result+'/' + args.save_name
    checkpoint = mx.callback.do_checkpoint(saveroot)


    model.fit(train,
              begin_epoch=0,
              num_epoch=100000,
              eval_data=val,
              eval_metric=Multi_Accuracy(num=3),#需要注意的地方
              optimizer='sgd',

              optimizer_params=optimizer_params,

              initializer=initializer,
              allow_missing=True,
              batch_end_callback=mx.callback.Speedometer(64, 50),

              epoch_end_callback=checkpoint
              )