纪念第一次用 Caffe

课程设计要拿 Caffe 搞一点事情。这完全就是搞事情嘛!时间辣么短……

需求

最终需求是使用 CaffeNet,将最后的几层给修改掉,做一下 Fine Tuning,然后使用最后一层来做特征提取。

降级的需求

入门嘛…… 先能 “用起来” 再说。

加载 Mnist 数据集训练好的模型,进行一些分类,并拿出来最后一层的数据当作特征。

prototxt 文件

lenet_train_test.prototxt 进行了少许修改。

首先,最上面去掉了 mnist 层,改成了下面的形式

name: "LeNet"
input:"data"
input_dim:1
input_dim:1
input_dim:28
input_dim:28

然后,最下面去掉 accuracyloss 层,改成一个 SoftMax 层:

layer {
  name: "answer"
  type: "Softmax"
  bottom: "ip2"
  top: "answer"
}

其余都不动。所以,最终的 prototxt 文件是这样的:

name: "LeNet"
input:"data"
input_dim:1
input_dim:1
input_dim:28
input_dim:28

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: 10
    weight_filler {
      type: "xavier"
    }
    bias_filler {
      type: "constant"
    }
  }
}
layer {
  name: "answer"
  type: "Softmax"
  bottom: "ip2"
  top: "answer"
}

代码

import sys
import os

# 由于设置关系,Pycharm 里面只好使用这种方式加载 Caffe
caffe_root = "/home/haoyu/ProgramFiles/nlpcaffe"
caffe_root_example = os.path.join(caffe_root, 'examples')
sys.path.insert(0, os.path.join(caffe_root, 'python'))
import caffe

# deploy 文件
file_deploy_proto = os.path.join(caffe_root_example, 'mnist/lenet_train_test_deploy.prototxt')
# 训练好的 caffemodel
file_caffe_model = os.path.join(caffe_root_example, 'mnist/lenet_iter_10000.caffemodel')
# 随便从 Mnist 数据集里面挑一张图片
img = '/home/haoyu/00015.png'

# 加载 model 和 network
net = caffe.Net(file_deploy_proto, file_caffe_model, caffe.TEST)

# 图片预处理设置
# 设定图片的 shape 格式 (1,1,28,28)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
# 改变维度的顺序,由原始图片 (28,28,1) 变为 (1,28,28)
transformer.set_transpose('data', (2, 0, 1))
# 缩放到 (0,255) 之间
transformer.set_raw_scale('data', 255)

# 使用灰度加载图片
im = caffe.io.load_image(img, color=False)
# 执行上面设置的图片预处理操作,并将图片载入到 blob 中
net.blobs['data'].data[...] = transformer.preprocess('data', im)

# 执行测试
out = net.forward()
# 取出最后一层(Softmax)属于某个类别的概率值
prob = net.blobs['answer'].data[0].flatten()
print(prob)
print('Answer:', prob.argsort()[-1])

# 查看各层数据规模
print([(k, v.data.shape) for k, v in net.blobs.items()])
# 提取某层数据(特征)
feature = net.blobs['ip2'].data
print(feature)

参考资料

留下评论