caffe的python接口使用方法示例

训练

import sys
#添加caffe下的python文件夹的路径
sys.path.append("../caffe-master/python")
import caffe
import numpy as np
import os

#使用caffemodel进行训练
weights = 'googlenet_bn_stepsize_6400_iter_1200000.caffemodel'
caffe.set_device(0)  #选择0号GPU
caffe.set_mode_gpu() #使用GPU模式
solver = caffe.SGDSolver('solver.prototxt') #加载solver文件
solver.net.copy_from(weights) #加载caffemodel
solver.solve() #开始运行

#使用solverstate继续之前的训练
s = "googlenet_bn_stepsize_6400_iter_1200000.solverstate"
caffe.set_device(0)  #选择0号GPU
caffe.set_mode_gpu() #使用GPU模式
solver = caffe.SGDSolver('solver.prototxt')#加载solver文件
solver.restore(s) #加载solverstate
solver.solve() #开始运行

提取特征

#这里是将特征写入txt，写入其他文件（如npy等）方法类似
import sys
#添加caffe下的python文件夹的路径
sys.path.append("../caffe-master/python")
import caffe
import os
import numpy as np

deploy = 'train_val.prototxt'  #CNN配置文件
caffe_model = 'googlenet_bn_stepsize_6400_iter_1200000.caffemodel' #权重文件

caffe.set_device(0)  #选择0号GPU
caffe.set_mode_gpu() #使用GPU模式
net = caffe.Net(deploy,caffe_model,caffe.TEST) #初始化网络，选取test模式，所以prototxt的输入层

#打开一个文本文件
file_name = 'fea.txt'
output = open(file_name,'w')

num = 100  #需要提取特征的图片的数量，这里是因为我在prototxt配置文件里面将batch size设置为1， 所以num就是输入层的图片list中图片的数量
while num:
    print num
    num = num - 1
    net.forward() #网络前传一次，后传类似

    param = net.params['conv1'][0].data #获取一层的权重
    fea = net.blobs['pool0'].data #获取一层的数据
    print np.shape(fea)  #输出数据的形状，这里是一个pooling层，所以应该是4维的
	
    #将特征连成一个字符串，写入文本文件
    str_fea = ''
    for x in xrange(len(fea[0])):
        str_fea = str_fea + str(fea[0][x][0][0]) +' '
    str_fea = str_fea + '\n'
    output.write(str_fea)
	
output.close()

使用输入图片到deploy

deploy中数据层：

input: "data"
input_dim: 1
input_dim: 3
input_dim: 512
input_dim: 512

将图片输入deploy：

import sys
sys.path.append("../caffe/python/")
import caffe
import os
import numpy as np
import cv2

deploy = 'deploy.prototxt'
caffe_model = 'googlenet_bn_stepsize_6400_iter_120000.caffemodel'

caffe.set_device(0)
caffe.set_mode_gpu()
net = caffe.Net(deploy,caffe_model,caffe.TEST)

im = cv2.imread('exmaple.jpg')
im = cv2.resize(im,(512,512))
im = np.array(im, np.float32)
im -= mean #减去每个通道的均值
im = np.transpose(im, (2, 0, 1)) #由图片的shape：512*512*3 变为caffe输入要求的3*512*512
net.blobs['data'].data[...] = im #将图片输入到网络
net.forward() #完成一次前传