在Caffe中實現模型融合

模型融合

有的時候我們手頭可能有了若干個已經訓練好的模型，這些模型可能是同樣的結構，也可能是不同的結構，訓練模型的數據可能是同一批，也可能不同。無論是出於要通過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