caffe框架学习六——Editing model parameters

英文原文 https://nbviewer.jupyter.org/github/BVLC/caffe/blob/master/examples/net_surgery.ipynb

任务 How to do net surgery and manually change model parameters for custom use.

# -*- coding: UTF-8 -*-
import numpy as np
import matplotlib.pyplot as plt

caffe_root = '/home/zjhao/caffe/'
import sys
sys.path.insert(0, caffe_root + 'python')

import caffe

plt.rcParams['figure.figsize'] = (10, 10)
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

# 设计filter
# 加载网络
caffe.set_mode_gpu()
net = caffe.Net('/home/zjhao/caffe/examples/net_surgery/conv.prototxt',
                caffe.TEST)
# 列出数据和参数
print 'blobs {}\n params {}'.format(net.blobs.keys(), net.params.keys())

# 加载一张样例图片
im = np.array(
    caffe.io.load_image(
        '/home/zjhao/desktop/caffe/cat_gray.jpg',
        color=False)).squeeze()  #从数组的形状中删除单维度条目，即把shape中为1的维度去掉
plt.title('original image')
plt.imshow(im)
plt.axis('off')
# plt.show()

im_input = im[np.newaxis, np.newaxis, :, :]
net.blobs['data'].reshape(*im_input.shape)
net.blobs['data'].data[...] = im_input


# filter输出的辅助函数
def show_filter(net):
    net.forward()
    plt.figure()
    filt_min = net.blobs['conv'].data.min()
    filt_max = net.blobs['conv'].data.max()
    for i in range(3):
        plt.subplot(1, 4, i + 2)
        plt.title('filter #{} output'.format(i))
        plt.imshow(net.blobs['conv'].data[0, i], vmin=filt_min, vmax=filt_max)
        plt.tight_layout()  #自动调整子图参数,使之填充整个图像.
        plt.axis('off')
    # plt.show()


# 使用初始值过滤图像
show_filter(net)

# 使用filter的偏置值相应地提高其输出
# 选择第一个filter的输出
conv0 = net.blobs['conv'].data[0, 0]
print 'pre-originalsurgery output mean {:.2f}'.format(conv0.mean())
# 设置第一original个filter的偏置为1
net.params['conv'][1].data[0] = 1
net.forward()
print 'post-surgery output mean {:.2f}'.format(conv0.mean())

ksize = net.params['conv'][0].data.shape[2:]
# 使用高斯模糊
sigma = 1
y, x = np.mgrid[-ksize[0] // 2 + 1:ksize[0] // 2 + 1, -ksize[1] // 2 +
                1:ksize[1] // 2 + 1]
g = np.exp(-(x**2 + y**2) / (2.0 * sigma**2))
gaussian = (g / g.sum()).astype(np.float32)
net.params['conv'][0].data[0] = gaussian
# 利用Sobel算子进行边缘检测
net.params['conv'][0].data[1:] = 0
sobel = np.array((-1, -2, -1, 0, 0, 0, 1, 2, 1), dtype=np.float32).reshape((3,
                                                                            3))
net.params['conv'][0].data[1, 0, 1:-1, 1:-1] = sobel  #水平的
net.params['conv'][0].data[2, 0, 1:-1, 1:-1] = sobel.T  #垂直的
show_filter(net)

# 加载原始的网络,并移开全连接层的参数
net = caffe.Net(
    caffe_root + 'models/bvlc_reference_caffenet/deploy.prototxt', caffe_root +
    'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel',
    caffe.TEST)
params = ['fc6', 'fc7', 'fc8']
# fc_params = {name: (weights, biases)}
fc_params = {
    pr: (net.params[pr][0].data, net.params[pr][1].data)
    for pr in params
}

for fc in params:
    print '{} weights are {} dimensinal and biases are {} dimensional'.format(
        fc, fc_params[fc][0].shape, fc_params[fc][1].shape)

# 加载完全卷积网络来移植参数
net_full_conv = caffe.Net(
    '/home/zjhao/caffe/examples/net_surgery/bvlc_caffenet_full_conv.prototxt',
    caffe_root +
    'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel',
    caffe.TEST)
params_full_conv = ['fc6-conv', 'fc7-conv', 'fc8-conv']
# conv_params = {name: (weights, biases)}
conv_params = {
    pr: (net_full_conv.params[pr][0].data, net_full_conv.params[pr][1].data)
    for pr in params_full_conv
}
for conv in params_full_conv:
    print '{} weights are {} dimensinal and biases are {} dimensional'.format(
        conv, conv_params[conv][0].shape, conv_params[conv][1].shape)

for pr, pr_conv in zip(params, params_full_conv):
    conv_params[pr_conv][0].flat = fc_params[pr][0].flat
    conv_params[pr_conv][1][...] = fc_params[pr][1]

net_full_conv.save('net_surgery/bvlc_caffenet_full_conv.caffemodel')

# 加载输入并配置预处理
im = caffe.io.load_image('/home/zjhao/desktop/caffe/cat_gray.jpg')
transformer = caffe.io.Transformer(
    {'data': net_full_conv.blobs['data'].data.shape})
transformer.set_mean(
    'data',
    np.load(caffe_root +
            'python/caffe/imagenet/ilsvrc_2012_mean.npy').mean(1).mean(1))
transformer.set_transpose('data', (2, 0, 1))
transformer.set_channel_swap('data', (2, 1, 0))
transformer.set_raw_scale('data', 255.0)
# 通过前向绘制分类图，并在每个位置打印预测指标
out = net_full_conv.forward_all(
    data=np.asarray([transformer.preprocess('data', im)]))
print out['prob'][0].argmax(axis=0)
# 显示净输入和置信度图(每个位置顶部预测的概率)
plt.subplot(1, 2, 1)
plt.imshow(transformer.deprocess('data', net_full_conv.blobs['data'].data[0]))
plt.subplot(1, 2, 2)
plt.imshow(out['prob'][0, 281])

# plt.show()