模型融合

有的时候我们手头可能有了若干个已经训练好的模型，这些模型可能是同样的结构，也可能是不同的结构，训练模型的数据可能是同一批，也可能不同。无论是出于要通过ensemble提升性能的目的，还是要设计特殊作用的网络，在用Caffe做工程时，融合都是一个常见的步骤。

比如考虑下面的场景，我们有两个模型，都是基于resnet-101，分别在两拨数据上训练出来的。我们希望把这两个模型的倒数第二层拿出来，接一个fc层然后训练这个fc层进行融合。那么有两个问题需要解决：1）两个模型中的层的名字都是相同的，但是不同模型对应的权重不同；2）如何同时在一个融合好的模型中把两个训练好的权重都读取进来。

Caffe中并没有直接用于融合的官方工具，本文介绍一个简单有效的土办法，用融合模型进行ensemble的例子，一步步实现模型融合。

完整例子

模型定义和脚本：

https://github.com/frombeijingwithlove/dlcv_for_beginners/tree/master/random_bonus/multiple_models_fusion_caffe

预训练模型：

https://github.com/frombeijingwithlove/dlcv_book_pretrained_caffe_models/blob/master/mnist_lenet_odd_iter_30000.caffemodel

https://github.com/frombeijingwithlove/dlcv_book_pretrained_caffe_models/blob/master/mnist_lenet_even_iter_30000.caffemodel

虽然模型只是简单的LeNet-5，但是方法是可以拓展到其他大模型上的。

模型(及数据)准备：直接采用预训练好的模型

本文的例子要融合的是两个不同任务的模型：

对偶数0, 2, 4, 6, 8分类的模型

对奇数1, 3, 5, 7, 9分类的模型

采用的网络都是LeNet-5

直接从上节中提到的本文例子的repo下载预定义的模型和权重。

上一部分第一个链接中已经写好了用来训练的LeNet-5结构和solver，用的是ImageData层，以训练奇数分类的模型为例：

name: "LeNet"

layer {

  name: "mnist"

  type: "ImageData"

  top: "data"

  top: "label"

  include {

    phase: TRAIN

  }

  transform_param {

    mean_value: 128

    scale: 0.00390625

  }

  image_data_param {

    source: "train_odd.txt"

    is_color: false

    batch_size: 25

  }

}

layer {

  name: "mnist"

  type: "ImageData"

  top: "data"

  top: "label"

  include {

    phase: TEST

  }

  transform_param {

    mean_value: 128

    scale: 0.00390625

  }

  image_data_param {

    source: "val_odd.txt"

    is_color: false

    batch_size: 20

  }

}

layer {

  name: "conv1"

  type: "Convolution"

  bottom: "data"

  top: "conv1"

  param {

    lr_mult: 1

  }

  param {

    lr_mult: 2

  }

  convolution_param {

    num_output: 20

    kernel_size: 5

    stride: 1

    weight_filler {

      type: "xavier"

    }

    bias_filler {

      type: "constant"

    }

  }

}

layer {

  name: "pool1"

  type: "Pooling"

  bottom: "conv1"

  top: "pool1"

  pooling_param {

    pool: MAX

    kernel_size: 2

    stride: 2

  }

}

layer {

  name: "conv2"

  type: "Convolution"

  bottom: "pool1"

  top: "conv2"

  param {

    lr_mult: 1

  }

  param {

    lr_mult: 2

  }

  convolution_param {

    num_output: 50

    kernel_size: 5

    stride: 1

    weight_filler {

      type: "xavier"

    }

    bias_filler {

      type: "constant"

    }

  }

}

layer {

  name: "pool2"

  type: "Pooling"

  bottom: "conv2"

  top: "pool2"

  pooling_param {

    pool: MAX

    kernel_size: 2

    stride: 2

  }

}

layer {

  name: "ip1"

  type: "InnerProduct"

  bottom: "pool2"

  top: "ip1"

  param {

    lr_mult: 1

  }

  param {

    lr_mult: 2

  }

  inner_product_param {

    num_output: 500

    weight_filler {

      type: "xavier"

    }

    bias_filler {

      type: "constant"

    }

  }

}

layer {

  name: "relu1"

  type: "ReLU"

  bottom: "ip1"

  top: "ip1"

}

layer {

  name: "ip2"

  type: "InnerProduct"

  bottom: "ip1"

  top: "ip2"

  param {

    lr_mult: 1

  }

  param {

    lr_mult: 2

  }

  inner_product_param {

    num_output: 5

    weight_filler {

      type: "xavier"

    }

    bias_filler {

      type: "constant"

    }

  }

}

layer {

  name: "accuracy"

  type: "Accuracy"

  bottom: "ip2"

  bottom: "label"

  top: "accuracy"

  include {

    phase: TEST

  }

}

layer {

  name: "loss"

  type: "SoftmaxWithLoss"

  bottom: "ip2"

  bottom: "label"

  top: "loss"

}

训练偶数分类的prototxt的唯一区别就是ImageData层中数据的来源不一样。

模型(及数据)准备：Start From Scratch

当然也可以自行训练这两个模型，毕竟只是个用于演示的小例子，很简单。方法如下：

第一步 下载MNIST数据

直接运行download_mnist.sh这个脚本

第二步 转换MNIST数据为图片

运行convert_mnist.py，可以从mnist.pkl.gz中提取所有图片为jpg

import os

import pickle, gzip

from matplotlib import pyplot

# Load the dataset

print('Loading data from mnist.pkl.gz ...')

with gzip.open('mnist.pkl.gz', 'rb') as f:

    train_set, valid_set, test_set = pickle.load(f)

imgs_dir = 'mnist'

os.system('mkdir -p {}'.format(imgs_dir))

datasets = {'train': train_set, 'val': valid_set, 'test': test_set}

for dataname, dataset in datasets.items():

    print('Converting {} dataset ...'.format(dataname))

    data_dir = os.sep.join([imgs_dir, dataname])

    os.system('mkdir -p {}'.format(data_dir))

    for i, (img, label) in enumerate(zip(*dataset)):

        filename = '{:0>6d}_{}.jpg'.format(i, label)

        filepath = os.sep.join([data_dir, filename])

        img = img.reshape((28, 28))

        pyplot.imsave(filepath, img, cmap='gray')

        if (i+1) % 10000 == 0:

            print('{} images converted!'.format(i+1))

第三步 生成奇数、偶数和全部数据的列表

运行gen_img_list.py，可以分别生成奇数、偶数和全部数据的训练及验证列表：

import os

import sys

mnist_path = 'mnist'

data_sets = ['train', 'val']

for data_set in data_sets:

    odd_list = '{}_odd.txt'.format(data_set)

    even_list = '{}_even.txt'.format(data_set)

    all_list = '{}_all.txt'.format(data_set)

    root = os.sep.join([mnist_path, data_set])

    filenames = os.listdir(root)

    with open(odd_list, 'w') as f_odd, open(even_list, 'w') as f_even, open(all_list, 'w') as f_all:

        for filename in filenames:

            filepath = os.sep.join([root, filename])

            label = int(filename[:filename.rfind('.')].split('_')[1])

            line = '{} {}\n'.format(filepath, label)

            f_all.write(line)

            line = '{} {}\n'.format(filepath, int(label/2))

            if label % 2:

                f_odd.write(line)

            else:

                f_even.write(line)

第四步 训练两个不同的模型

就直接训练就行了。Solver的例子如下：

net: "lenet_odd_train_val.prototxt"

test_iter: 253

test_initialization: false

test_interval: 1000

base_lr: 0.01

momentum: 0.9

weight_decay: 0.0005

lr_policy: "step"

gamma: 0.707

stepsize: 1000

display: 200

max_iter: 30000

snapshot: 30000

snapshot_prefix: "mnist_lenet_odd"

solver_mode: GPU

注意到test_iter是个奇怪的253，这是因为MNIST的验证集中奇数样本多一些，一共是5060个，训练随便取个30个epoch，应该是够了。

制作融合后模型的网络定义

前面提到了模型融合的难题之一在于层的名字可能是相同的，解决这个问题非常简单，只要把名字改成不同就可以，加个前缀就行。按照这个思路，我们给奇数分类和偶数分类的模型的每层前分别加上odd/和even/作为前缀，同时我们给每层的学习率置为0，这样融合的时候就可以只训练融合的全连接层就可以了。实现就是用Python自带的正则表达式匹配，然后进行字符串替换，代码就是第一部分第一个链接中的rename_n_freeze_layers.py：

import sys

import re

layer_name_regex = re.compile('name:\s*"(.*?)"')

lr_mult_regex = re.compile('lr_mult:\s*\d+\.*\d*')

input_filepath = sys.argv[1]

output_filepath = sys.argv[2]

prefix = sys.argv[3]

with open(input_filepath, 'r') as fr, open(output_filepath, 'w') as fw:

    prototxt = fr.read()

    layer_names = set(layer_name_regex.findall(prototxt))

    for layer_name in layer_names:

        prototxt = prototxt.replace(layer_name, '{}/{}'.format(prefix, layer_name))

    lr_mult_statements = set(lr_mult_regex.findall(prototxt))

    for lr_mult_statement in lr_mult_statements:

        prototxt = prototxt.replace(lr_mult_statement, 'lr_mult: 0')

    fw.write(prototxt)

这个方法虽然土，不过有效，另外需要注意的是如果确定不需要动最后一层以外的参数，或者原始的训练prototxt中就没有lr_mult的话，可以考虑用Caffe的propagate_down这个参数。把这个脚本分别对奇数和偶数模型执行，并记住自己设定的前缀even和odd，然后把数据层到ip1层的定义复制并粘贴到一个文件中，然后把ImageData层和融合层的定义也写入到这个文件，注意融合前需要先用Concat层把特征拼接一下：

name: "LeNet"

layer {

  name: "mnist"

  type: "ImageData"

  top: "data"

  top: "label"

  include {

    phase: TRAIN

  }

  transform_param {

    mean_value: 128

    scale: 0.00390625

  }

  image_data_param {

    source: "train_all.txt"

    is_color: false

    batch_size: 50

  }

}

layer {

  name: "mnist"

  type: "ImageData"

  top: "data"

  top: "label"

  include {

    phase: TEST

  }

  transform_param {

    mean_value: 128

    scale: 0.00390625

  }

  image_data_param {

    source: "val_all.txt"

    is_color: false

    batch_size: 20

  }

}

...

### rename_n_freeze_layers.py 生成的网络结构部分 ###

...

layer {

  name: "concat"

  bottom: "odd/ip1"

  bottom: "even/ip1"

  top: "ip1_fused"

  type: "Concat"

  concat_param {

    axis: 1

  }

}

layer {

  name: "ip2"

  type: "InnerProduct"

  bottom: "ip1_fused"

  top: "ip2"

  param {

    lr_mult: 1

  }

  param {

    lr_mult: 2

  }

  inner_product_param {

    num_output: 10

    weight_filler {

      type: "xavier"

    }

    bias_filler {

      type: "constant"

    }

  }

}

layer {

  name: "accuracy"

  type: "Accuracy"

  bottom: "ip2"

  bottom: "label"

  top: "accuracy"

  include {

    phase: TEST

  }

}

layer {

  name: "loss"

  type: "SoftmaxWithLoss"

  bottom: "ip2"

  bottom: "label"

  top: "loss"

}

分别读取每个模型的权重并生成融合模型的权重

这个思路就是用pycaffe进行读取，然后按照层名字的对应关系进行值拷贝，最后再存一下就可以，代码如下：

import sys

sys.path.append('/path/to/caffe/python')

import caffe

fusion_net = caffe.Net('lenet_fusion_train_val.prototxt', caffe.TEST)

model_list = [

    ('even', 'lenet_even_train_val.prototxt', 'mnist_lenet_even_iter_30000.caffemodel'),

    ('odd', 'lenet_odd_train_val.prototxt', 'mnist_lenet_odd_iter_30000.caffemodel')

]

for prefix, model_def, model_weight in model_list:

    net = caffe.Net(model_def, model_weight, caffe.TEST)

    for layer_name, param in net.params.iteritems():

        n_params = len(param)

        try:

            for i in range(n_params):

                net.params['{}/{}'.format(prefix, layer_name)][i].data[...] = param[i].data[...]

        except Exception as e:

            print(e)

fusion_net.save('init_fusion.caffemodel')

训练融合后的模型

这个也没什么好说的了，直接训练即可，本文例子的参考Solver如下：

net: "lenet_fusion_train_val.prototxt"

test_iter: 500

test_initialization: false

test_interval: 1000

base_lr: 0.01

momentum: 0.9

weight_decay: 0.0005

lr_policy: "step"

gamma: 0.707

stepsize: 1000

display: 200

max_iter: 30000

snapshot: 30000

snapshot_prefix: "mnist_lenet_fused"

solver_mode: GPU