├── README.md ├── classification.py ├── detection.py ├── generate_lmdb.py ├── prototxt.py ├── prune.py ├── quantize.py ├── solver.py └── train_val.py /README.md: -------------------------------------------------------------------------------- 1 | # Caffe-Python-Tutorial 2 | -------------------------------------------------------------------------------- /classification.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 用于模型的单张图像分类操作 3 | import os 4 | os.environ['GLOG_minloglevel'] = '2' # 将caffe的输出log信息不显示，必须放到import caffe前 5 | import caffe # caffe 模块 6 | from caffe.proto import caffe_pb2 7 | from google.protobuf import text_format 8 | import numpy as np 9 | import cv2 10 | import matplotlib.pyplot as plt 11 | import time 12 | 13 | # 分类单张图像img 14 | def classification(img, net, transformer, synset_words): 15 | im = caffe.io.load_image(img) 16 | # 导入输入图像 17 | net.blobs['data'].data[...] = transformer.preprocess('data', im) 18 | 19 | start = time.clock() 20 | # 执行测试 21 | net.forward() 22 | end = time.clock() 23 | print('classification time: %f s' % (end - start)) 24 | 25 | # 查看目标检测结果 26 | labels = np.loadtxt(synset_words, str, delimiter='\t') 27 | 28 | category = net.blobs['prob'].data[0].argmax() 29 | 30 | class_str = labels[int(category)].split(',') 31 | class_name = class_str[0] 32 | # text_font = cv2.cv.InitFont(cv2.cv.CV_FONT_HERSHEY_SCRIPT_SIMPLEX, 1, 1, 0, 3, 8) 33 | cv2.putText(im, class_name, (0, im.shape[0]), cv2.cv.CV_FONT_HERSHEY_SIMPLEX, 1, (55, 255, 155), 2) 34 | 35 | # 显示结果 36 | plt.imshow(im, 'brg') 37 | plt.show() 38 | 39 | #CPU或GPU模型转换 40 | caffe.set_mode_cpu() 41 | #caffe.set_device(0) 42 | #caffe.set_mode_gpu() 43 | 44 | caffe_root = '../../' 45 | # 网络参数（权重）文件 46 | caffemodel = caffe_root + 'models/bvlc_alexnet/bvlc_alexnet.caffemodel' 47 | # 网络实施结构配置文件 48 | deploy = caffe_root + 'models/bvlc_alexnet/deploy.prototxt' 49 | 50 | 51 | img_root = caffe_root + 'data/VOCdevkit/VOC2007/JPEGImages/' 52 | synset_words = caffe_root + 'data/ilsvrc12/synset_words.txt' 53 | 54 | # 网络实施分类 55 | net = caffe.Net(deploy, # 定义模型结构 56 | caffemodel, # 包含了模型的训练权值 57 | caffe.TEST) # 使用测试模式(不执行dropout) 58 | 59 | # 加载ImageNet图像均值 (随着Caffe一起发布的) 60 | mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy') 61 | mu = mu.mean(1).mean(1) # 对所有像素值取平均以此获取BGR的均值像素值 62 | 63 | # 图像预处理 64 | transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape}) 65 | transformer.set_transpose('data', (2,0,1)) 66 | transformer.set_mean('data', mu) 67 | transformer.set_raw_scale('data', 255) 68 | transformer.set_channel_swap('data', (2,1,0)) 69 | 70 | # 处理图像 71 | while 1: 72 | img_num = raw_input("Enter Img Number: ") 73 | if img_num == '': break 74 | img = img_root + '{:0>6}'.format(img_num) + '.jpg' 75 | classification(img,net,transformer,synset_words) 76 | 77 | -------------------------------------------------------------------------------- /detection.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 用于模型的单张图像分类操作 3 | import os 4 | os.environ['GLOG_minloglevel'] = '2' # 将caffe的输出log信息不显示，必须放到import caffe前 5 | import caffe # caffe 模块 6 | from caffe.proto import caffe_pb2 7 | from google.protobuf import text_format 8 | import numpy as np 9 | import cv2 10 | import matplotlib.pyplot as plt 11 | import time 12 | 13 | # 分类单张图像img 14 | def detection(img, net, transformer, labels_file): 15 | im = caffe.io.load_image(img) 16 | net.blobs['data'].data[...] = transformer.preprocess('data', im) 17 | 18 | start = time.clock() 19 | # 执行测试 20 | net.forward() 21 | end = time.clock() 22 | print('detection time: %f s' % (end - start)) 23 | 24 | # 查看目标检测结果 25 | file = open(labels_file, 'r') 26 | labelmap = caffe_pb2.LabelMap() 27 | text_format.Merge(str(file.read()), labelmap) 28 | 29 | loc = net.blobs['detection_out'].data[0][0] 30 | confidence_threshold = 0.5 31 | for l in range(len(loc)): 32 | if loc[l][2] >= confidence_threshold: 33 | xmin = int(loc[l][3] * im.shape[1]) 34 | ymin = int(loc[l][4] * im.shape[0]) 35 | xmax = int(loc[l][5] * im.shape[1]) 36 | ymax = int(loc[l][6] * im.shape[0]) 37 | img = np.zeros((512, 512, 3), np.uint8) # 生成一个空彩色图像 38 | cv2.rectangle(im, (xmin, ymin), (xmax, ymax), (55 / 255.0, 255 / 255.0, 155 / 255.0), 2) 39 | 40 | # 确定分类类别 41 | class_name = labelmap.item[int(loc[l][1])].display_name 42 | # text_font = cv2.cv.InitFont(cv2.cv.CV_FONT_HERSHEY_SCRIPT_SIMPLEX, 1, 1, 0, 3, 8) 43 | cv2.putText(im, class_name, (xmin, ymax), cv2.cv.CV_FONT_HERSHEY_SIMPLEX, 1, (55, 255, 155), 2) 44 | 45 | # 显示结果 46 | plt.imshow(im, 'brg') 47 | plt.show() 48 | 49 | #CPU或GPU模型转换 50 | caffe.set_mode_cpu() 51 | #caffe.set_device(0) 52 | #caffe.set_mode_gpu() 53 | 54 | caffe_root = '../../' 55 | # 网络参数（权重）文件 56 | caffemodel = caffe_root + 'models/SSD_300x300/VGG_VOC0712_SSD_300x300_iter_60000.caffemodel' 57 | # 网络实施结构配置文件 58 | deploy = caffe_root + 'models/SSD_300x300/deploy.prototxt' 59 | 60 | 61 | img_root = caffe_root + 'data/VOCdevkit/VOC2007/JPEGImages/' 62 | labels_file = caffe_root + 'data/VOC0712/labelmap_voc.prototxt' 63 | 64 | # 网络实施分类 65 | net = caffe.Net(deploy, # 定义模型结构 66 | caffemodel, # 包含了模型的训练权值 67 | caffe.TEST) # 使用测试模式(不执行dropout) 68 | 69 | # 加载ImageNet图像均值 (随着Caffe一起发布的) 70 | mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy') 71 | mu = mu.mean(1).mean(1) # 对所有像素值取平均以此获取BGR的均值像素值 72 | 73 | # 图像预处理 74 | transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape}) 75 | transformer.set_transpose('data', (2,0,1)) 76 | transformer.set_mean('data', mu) 77 | transformer.set_raw_scale('data', 255) 78 | transformer.set_channel_swap('data', (2,1,0)) 79 | 80 | # 处理图像 81 | while 1: 82 | img_num = raw_input("Enter Img Number: ") 83 | if img_num == '': break 84 | img = img_root + '{:0>6}'.format(img_num) + '.jpg' 85 | detection(img,net,transformer,labels_file) -------------------------------------------------------------------------------- /generate_lmdb.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 将图像数据生成lmdb数据集 3 | # 1. 生成分类图像数据集 4 | # 2. 生成目标检测图像数据集 5 | import os 6 | import sys 7 | import numpy as np 8 | import random 9 | from caffe.proto import caffe_pb2 10 | from xml.dom.minidom import parse 11 | 12 | # 生成分类标签文件 13 | def labelmap(labelmap_file, label_info): 14 | labelmap = caffe_pb2.LabelMap() 15 | for i in range(len(label_info)): 16 | labelmapitem = caffe_pb2.LabelMapItem() 17 | labelmapitem.name = label_info[i]['name'] 18 | labelmapitem.label = label_info[i]['label'] 19 | labelmapitem.display_name = label_info[i]['display_name'] 20 | labelmap.item.add().MergeFrom(labelmapitem) 21 | with open(labelmap_file, 'w') as f: 22 | f.write(str(labelmap)) 23 | 24 | def rename_img(Img_dir): 25 | # 重新命名Img,这里假设图像名称表示为000011.jpg、003456.jpg、000000.jpg格式，最高6位，前补0 26 | # 列出图像，并将图像改为序号名称 27 | listfile=os.listdir(Img_dir) # 提取图像名称列表 28 | total_num = 0 29 | for line in listfile: #把目录下的文件都赋值给line这个参数 30 | if line[-4:] == '.jpg': 31 | newname = '{:0>6}'.format(total_num) +'.jpg' 32 | os.rename(os.path.join(Img_dir, line), os.path.join(Img_dir, newname)) 33 | total_num+=1 #统计所有图像 34 | 35 | def get_img_size(): 36 | pass 37 | 38 | def create_annoset(anno_args): 39 | if anno_args.anno_type == "detection": 40 | cmd = "E:\Code\windows-ssd/Build/x64/Release/convert_annoset.exe" \ 41 | " --anno_type={}" \ 42 | " --label_type={}" \ 43 | " --label_map_file={}" \ 44 | " --check_label={}" \ 45 | " --min_dim={}" \ 46 | " --max_dim={}" \ 47 | " --resize_height={}" \ 48 | " --resize_width={}" \ 49 | " --backend={}" \ 50 | " --shuffle={}" \ 51 | " --check_size={}" \ 52 | " --encode_type={}" \ 53 | " --encoded={}" \ 54 | " --gray={}" \ 55 | " {} {} {}" \ 56 | .format(anno_args.anno_type, anno_args.label_type, anno_args.label_map_file, anno_args.check_label, 57 | anno_args.min_dim, anno_args.max_dim, anno_args.resize_height, anno_args.resize_width, anno_args.backend, anno_args.shuffle, 58 | anno_args.check_size, anno_args.encode_type, anno_args.encoded, anno_args.gray, anno_args.root_dir, anno_args.list_file, anno_args.out_dir) 59 | elif anno_args.anno_type == "classification": 60 | cmd = "E:\Code\windows-ssd/Build/x64/Release/convert_annoset.exe" \ 61 | " --anno_type={}" \ 62 | " --min_dim={}" \ 63 | " --max_dim={}" \ 64 | " --resize_height={}" \ 65 | " --resize_width={}" \ 66 | " --backend={}" \ 67 | " --shuffle={}" \ 68 | " --check_size={}" \ 69 | " --encode_type={}" \ 70 | " --encoded={}" \ 71 | " --gray={}" \ 72 | " {} {} {}" \ 73 | .format(anno_args.anno_type, anno_args.min_dim, anno_args.max_dim, anno_args.resize_height, 74 | anno_args.resize_width, anno_args.backend, anno_args.shuffle, anno_args.check_size, anno_args.encode_type, anno_args.encoded, 75 | anno_args.gray, anno_args.root_dir, anno_args.list_file, anno_args.out_dir) 76 | print cmd 77 | os.system(cmd) 78 | 79 | def detection_list(Img_dir, Ano_dir, Data_dir, test_num): 80 | # 造成目标检测图像数据库 81 | # Img_dir表示图像文件夹 82 | # Ano_dir表示图像标记文件夹，用labelImg生成 83 | # Data_dir生成的数据库文件地址 84 | # test_num测试图像的数目 85 | # 列出图像 86 | listfile=os.listdir(Img_dir) # 提取图像名称列表 87 | 88 | # 列出图像，并将图像改为序号名称 89 | total_num = 0 90 | for line in listfile: #把目录下的文件都赋值给line这个参数 91 | if line[-4:] == '.jpg': 92 | total_num+=1 #统计所有图像 93 | 94 | trainval_num = total_num-test_num # 训练图像数目 95 | 96 | # 生成训练图像及测试图像列表 97 | test_list_file=open(Data_dir+'/test.txt','w') 98 | train_list_file=open(Data_dir+'/trainval.txt','w') 99 | 100 | test_list = np.random.randint(0,total_num-1, size=test_num) 101 | 102 | train_list = range(total_num) 103 | for n in range(test_num): 104 | train_list.remove(test_list[n]) 105 | random.shuffle(train_list) 106 | 107 | # 测试图像排序，而训练图像不用排序 108 | test_list = np.sort(test_list) 109 | # train_list = np.sort(train_list) 110 | 111 | for n in range(trainval_num): 112 | train_list_file.write(Img_dir + '{:0>6}'.format(train_list[n]) +'.jpg '+ Ano_dir + '{:0>6}'.format(train_list[n]) +'.xml\n') 113 | 114 | for n in range(test_num): 115 | test_list_file.write(Img_dir + '{:0>6}'.format(test_list[n]) +'.jpg '+ Ano_dir + '{:0>6}'.format(test_list[n]) +'.xml\n') 116 | 117 | 118 | caffe_root = 'E:/Code/Github/windows_caffe/' 119 | data_root = caffe_root + 'data/mnist/' 120 | Img_dir = data_root + 'JPEGImages/' 121 | Ano_dir = data_root + 'Annotations/' 122 | anno_type = "detection" 123 | test_num = 100 124 | 125 | # 第一步，预处理图像，重命名图像名，生成各图像标记信息 126 | # rename_img(Img_dir) 127 | # 然后通过labelImg(可以通过pip install labelImg安装，出现错误可以删除PyQt4的描述）来生成图像的标记 128 | 129 | # 第二步，生成分类标签文件 130 | # 编辑label信息 131 | label_info = [ 132 | dict(name='none', label=0, display_name='background'), # 背景 133 | dict(name="cat",label=1, display_name='cat'), # 背景 134 | dict(name="dog",label=2, display_name='dog'), # 背景 135 | ] 136 | labelmap(data_root+'labelmap_voc.prototxt', label_info) 137 | 138 | # 第三步，生成图像及标记的列表文件 139 | if anno_type == "detection": 140 | detection_list(Img_dir, Ano_dir, data_root, test_num) 141 | else: 142 | # 分类，生成 143 | pass 144 | 145 | # 第四步，生成lmdb文件 146 | # 初始化信息 147 | anno_args = {} 148 | anno_args['anno_type'] = anno_type 149 | # 仅用于目标检测，lable文件的类型：{xml, json, txt} 150 | anno_args['label_type'] = "xml" 151 | # 仅用于目标检测，label文件地址 152 | anno_args['label_map_file'] = data_root+"labelmap_voc.prototxt" 153 | # 是否检测所有数据有相同的大小.默认False 154 | anno_args['check_size'] = False 155 | # 检测label是否相同的名称，默认False 156 | anno_args['check_label'] = False 157 | # 为0表示图像不用重新调整尺寸 158 | anno_args['min_dim'] = 0 159 | anno_args['max_dim'] = 0 160 | anno_args['resize_height'] = 0 161 | anno_args['resize_width'] = 0 162 | anno_args['backend'] = "lmdb" # 数据集格式（lmdb, leveldb） 163 | anno_args['shuffle'] = False # 是否随机打乱图像及对应标签 164 | anno_args['encode_type'] = "" # 图像编码格式('png','jpg',...) 165 | anno_args['encoded'] = False # 是否编码，默认False 166 | anno_args['gray'] = False # 是否视为灰度图，默认False 167 | anno_args['root_dir'] = data_root # 存放图像文件夹及标签文件夹的根目录 168 | anno_args['list_file'] = data_root + '' # listfile文件地址 169 | anno_args['out_dir'] = data_root # 最终lmdb的存在地址 170 | 171 | # 生成训练数据集train_lmdb 172 | anno_args['list_file'] = data_root + 'trainval.txt' 173 | create_annoset(anno_args) 174 | 175 | # 生成测试数据集train_lmdb 176 | anno_args['list_file'] = data_root + 'test.txt' 177 | create_annoset(anno_args) 178 | 179 | 180 | 181 | 182 | 183 | -------------------------------------------------------------------------------- /prototxt.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 生成Minst网络结构文件train.prototxt、test.prototxt及deploy.prototxt 3 | from __future__ import print_function 4 | import caffe 5 | from caffe import layers as L 6 | from caffe import params as P 7 | 8 | # 此函数生成LeNet5的主体结构 9 | def lenet5_body(net, from_layer): 10 | # 网络参数 11 | kwargs = { 12 | # param定义学习率，这里是指基础学习率step的情况，lt_mult乘以基础学习率为实际学习率，为0表示权重不更新，decay_mult同权重衰减相关 13 | 'param': [dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)], 14 | 'weight_filler': dict(type='xavier'), # 权重初始化模式 15 | 'bias_filler': dict(type='constant', value=0)} # 权重偏差初始化模式 16 | 17 | # 判断是否存在from_layer层 18 | assert from_layer in net.keys() 19 | # conv1 20 | net.conv1 = L.Convolution(net[from_layer], kernel_size=5, stride=1, num_output=20, pad=0, **kwargs) 21 | net.pool1 = L.Pooling(net.conv1, pool=P.Pooling.MAX, kernel_size=2, stride=2) 22 | net.conv2 = L.Convolution(net.pool1, kernel_size=5, stride=1, num_output=50, pad=0, **kwargs) 23 | net.pool2 = L.Pooling(net.conv2, pool=P.Pooling.MAX, kernel_size=2, stride=2) 24 | net.ip1 = L.InnerProduct(net.pool2, num_output=500, **kwargs) 25 | net.relu1 = L.ReLU(net.ip1, in_place=True) 26 | net.ip2 = L.InnerProduct(net.relu1, name='ip2', num_output=10, **kwargs) 27 | 28 | caffe_root = '../../' 29 | # caffe_root = 'E:/Code/Github/windows_caffe/' 30 | model_root = caffe_root + 'models/mnist/' 31 | 32 | # 训练数据 33 | train_data = caffe_root + "data/mnist/mnist_train_lmdb" 34 | # 测试数据 35 | test_data = caffe_root + "data/mnist/mnist_test_lmdb" 36 | 37 | # 训练网络 38 | train_net = caffe.NetSpec() # 基础网络 39 | # 带标签的数据输入层 40 | train_net.data, train_net.label = L.Data(source=train_data,backend=P.Data.LMDB, batch_size=64,ntop=2,transform_param=dict(scale=0.00390625)) 41 | # 生成LeNet5的主体结构 42 | lenet5_body(train_net, 'data') 43 | # 生成误差损失层 44 | train_net.loss = L.SoftmaxWithLoss(train_net.ip2, train_net.label) 45 | 46 | # 测试网络 47 | test_net = caffe.NetSpec() # 基础网络 48 | # 带标签的数据输入层 49 | test_net.data, test_net.label = L.Data(source=test_data, batch_size=100, backend=P.Data.LMDB, ntop=2,transform_param=dict(scale=0.00390625)) 50 | # 生成LeNet5的主体结构 51 | lenet5_body(test_net, 'data') 52 | # 生成误差损失层 53 | test_net.loss = L.SoftmaxWithLoss(test_net.ip2, test_net.label) 54 | # 添加一个精确层 55 | test_net.accuracy = L.Accuracy(test_net.ip2, test_net.label) 56 | 57 | # 实施网络 58 | deploy_net = caffe.NetSpec() # 基础网络 59 | # 带标签的数据输入层 60 | deploy_net.data = L.Input(input_param=dict(shape=dict(dim=[64,1,28,28]))) 61 | # 生成LeNet5的主体结构 62 | lenet5_body(deploy_net, 'data') 63 | deploy_net.prob = L.Softmax(deploy_net.ip2) 64 | 65 | # 保存训练文件 66 | with open(model_root+'train.prototxt', 'w') as f: 67 | print('name: "LenNet5_train"', file=f) 68 | print(train_net.to_proto(), file=f) 69 | 70 | with open(model_root+'test.prototxt', 'w') as f: 71 | print('name: "LenNet5_test"', file=f) 72 | print(test_net.to_proto(), file=f) 73 | 74 | with open(model_root+'deploy.prototxt', 'w') as f: 75 | print('name: "LenNet5_test"', file=f) 76 | print(deploy_net.to_proto(), file=f) -------------------------------------------------------------------------------- /prune.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 用于修剪网络模型 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | import os 6 | os.environ['GLOG_minloglevel'] = '2' 7 | import caffe 8 | 9 | # 由稠密变成CSC稀疏矩阵 10 | def dense_to_sparse_csc(W_flatten, num_level): 11 | # W_flatten: 扁平化的权重矩阵 12 | # num_level: 量化级别 13 | csc_W = [] # 存储稀疏矩阵 14 | csc_indx = [] 15 | indx = 0 16 | for n in range(len(W_flatten)): 17 | if W_flatten[n]!=0 or indx == 2**num_level: 18 | csc_W.append(W_flatten[n]) 19 | csc_indx.append(indx) 20 | indx = 0 21 | else: 22 | indx += 1 23 | if indx!=0: 24 | csc_W.append(0.0) 25 | csc_indx.append(indx-1) 26 | return np.array(csc_W, dtype=np.float32),np.array(csc_indx, dtype=np.int8) 27 | 28 | # 由稠密变成CSC稀疏矩阵 29 | def sparse_to_dense_csc(csc_W, csc_W_indx): 30 | # W_flatten: 扁平化的权重矩阵 31 | # num_level: 量化级别 32 | W_flatten = [] # 存储稠密矩阵 33 | indx = 0 34 | for n in range(len(csc_W)): 35 | if csc_W_indx[n]!=0: 36 | W_flatten.extend([0]*(csc_W_indx[n])) 37 | W_flatten.append(csc_W[n]) 38 | return np.array(W_flatten, dtype=np.float32) 39 | 40 | 41 | def read_sparse_net(filename, net, layers): 42 | pass 43 | 44 | def write_sparse_net(filename, net): 45 | pass 46 | 47 | # 画出各层参数的直方图 48 | def draw_hist_weight(net, layers): 49 | plt.figure() # 画图 50 | layer_num = len(layers) 51 | for i, layer in enumerate(layers): 52 | i += 1 53 | W = net.params[layer][0].data 54 | 55 | plt.subplot(layer_num/2, 2, i) 56 | numBins = 2 ^ 5 57 | plt.hist(W.flatten(), numBins, color='blue', alpha=0.8) 58 | plt.title(layer) 59 | plt.show() 60 | 61 | # 网络模型的参数 62 | def analyze_param(net, layers): 63 | 64 | print '\n=============analyze_param start===============' 65 | total_nonzero = 0 66 | total_allparam = 0 67 | percentage_list = [] 68 | for i, layer in enumerate(layers): 69 | i += 1 70 | W = net.params[layer][0].data 71 | b = net.params[layer][1].data 72 | 73 | print 'W(%s) range = [%f, %f]' % (layer, min(W.flatten()), max(W.flatten())) 74 | print 'W(%s) mean = %f, std = %f' % (layer, np.mean(W.flatten()), np.std(W.flatten())) 75 | non_zero = (np.count_nonzero(W.flatten()) + np.count_nonzero(b.flatten())) # 参数非零值 76 | all_param = (np.prod(W.shape) + np.prod(b.shape)) # 所有参数的数目 77 | this_layer_percentage = non_zero / float(all_param) # 参数比例 78 | total_nonzero += non_zero 79 | total_allparam += all_param 80 | print 'non-zero W and b cnt = %d' % non_zero 81 | print 'total W and b cnt = %d' % all_param 82 | print 'percentage = %f\n' % (this_layer_percentage) 83 | percentage_list.append(this_layer_percentage) 84 | 85 | print '=====> summary:' 86 | print 'non-zero W and b cnt = %d' % total_nonzero 87 | print 'total W and b cnt = %d' % total_allparam 88 | print 'percentage = %f' % (total_nonzero / float(total_allparam)) 89 | print '=============analyze_param ends ===============' 90 | return (total_nonzero / float(total_allparam), percentage_list) 91 | 92 | def prune(threshold, test_net, layers): 93 | sqarse_net = {} 94 | 95 | for i, layer in enumerate(layers): 96 | 97 | print '\n============ Pruning %s : threshold=%0.2f ============' % (layer,threshold[i]) 98 | W = test_net.params[layer][0].data 99 | b = test_net.params[layer][1].data 100 | hi = np.max(np.abs(W.flatten())) 101 | hi = np.sort(-np.abs(W.flatten()))[int((len(W.flatten())-1)* threshold[i])] 102 | 103 | # abs(val) = 0 ==> 0 104 | # abs(val) >= threshold ==> 1 105 | interpolated = np.interp(np.abs(W), [0, hi * threshold[i], 999999999.0], [0.0, 1.0, 1.0]) 106 | 107 | # 小于阈值的权重被随机修剪 108 | random_samps = np.random.rand(len(W.flatten())) 109 | random_samps.shape = W.shape 110 | 111 | # 修剪阈值 112 | # mask = (random_samps < interpolated) 113 | mask = (np.abs(W) > (np.abs(hi))) 114 | mask = np.bool_(mask) 115 | W = W * mask 116 | 117 | print 'non-zero W percentage = %0.5f ' % (np.count_nonzero(W.flatten()) / float(np.prod(W.shape))) 118 | # 保存修剪后的阈值 119 | test_net.params[layer][0].data[...] = W 120 | # net.params[layer][0].mask[...] = mask 121 | csc_W, csc_W_indx = dense_to_sparse_csc(W.flatten(), 8) 122 | dense_W = sparse_to_dense_csc(csc_W, csc_W_indx) 123 | sqarse_net[layer + '_W'] = csc_W 124 | sqarse_net[layer + '_W_indx'] = csc_W_indx 125 | 126 | # 计算修剪后的权重稀疏度 127 | # np.savez(model_dir + model_name +"_crc.npz",sqarse_net) # 保存存储成CRC格式的稀疏网络 128 | (total_percentage, percentage_list) = analyze_param(test_net, layers) 129 | test_loss, accuracy = test_net_accuracy(test_net) 130 | return (threshold, total_percentage, percentage_list, test_loss, accuracy) 131 | 132 | def test_net_accuracy(test_net): 133 | test_iter = 100 134 | test_loss = 0 135 | accuracy = 0 136 | for test_it in range(test_iter): 137 | # 进行一次测试 138 | test_net.forward() 139 | # 计算test loss 140 | test_loss += test_net.blobs['loss'].data 141 | # 计算test accuracy 142 | accuracy += test_net.blobs['accuracy'].data 143 | 144 | return (test_loss / test_iter), (accuracy / test_iter) 145 | 146 | 147 | def eval_prune_threshold(threshold_list, test_prototxt, caffemodel, prune_layers): 148 | def net_prune(threshold, test_prototx, caffemodel, prune_layers): 149 | test_net = caffe.Net(test_prototx, caffemodel, caffe.TEST) 150 | return prune(threshold, test_net, prune_layers) 151 | 152 | accuracy = [] 153 | for threshold in threshold_list: 154 | results = net_prune(threshold, test_prototxt, caffemodel, prune_layers) 155 | print 'threshold: ', results[0] 156 | print '\ntotal_percentage: ', results[1] 157 | print '\npercentage_list: ', results[2] 158 | print '\ntest_loss: ', results[3] 159 | print '\naccuracy: ', results[4] 160 | accuracy.append(results[4]) 161 | plt.plot(accuracy,'r.') 162 | plt.show() 163 | 164 | # 迭代训练修剪后网络 165 | def retrain_pruned(solver, pruned_caffemodel, threshold, prune_layers): 166 | #solver = caffe.SGDSolver(solver_proto) 167 | retrain_iter = 20 168 | 169 | accuracys = [] 170 | for i in range(retrain_iter): 171 | solver.net.copy_from(pruned_caffemodel) 172 | # solver.solve() 173 | solver.step(500) 174 | _,_,_,_,accuracy=prune(threshold, solver.test_nets[0], prune_layers) 175 | solver.test_nets[0].save(pruned_caffemodel) 176 | accuracys.append(accuracy) 177 | 178 | plt.plot(accuracys, 'r.-') 179 | plt.show() 180 | 181 | 182 | #CPU或GPU模型转换 183 | #caffe.set_mode_cpu() 184 | caffe.set_device(0) 185 | caffe.set_mode_gpu() 186 | 187 | caffe_root = '../../' 188 | #model_dir = caffe_root + 'models/SSD_300x300/' 189 | #deploy = model_dir + 'deploy.prototxt' 190 | #model_name = 'VGG_VOC0712_SSD_300x300_iter_60000' 191 | #caffemodel = model_dir + model_name + '.caffemodel' 192 | 193 | model_dir = caffe_root + 'models/mnist/' 194 | deploy = model_dir + 'deploy.prototxt' 195 | model_name = 'LeNet5_Mnist_shapshot_iter_10000' 196 | caffemodel = model_dir + model_name + '.caffemodel' 197 | test_prototxt = model_dir + 'test.prototxt' 198 | solver_proto = model_dir + 'solver.prototxt' 199 | 200 | solver = caffe.SGDSolver(solver_proto) 201 | 202 | # 要修剪的层 203 | prune_layers = ['conv1','conv2','ip1','ip2'] 204 | # 测试修剪率 205 | test_threshold_list = [[0.3, 1 ,1 ,1], [0.4, 1 ,1 ,1], [0.5, 1 ,1 ,1], [0.6, 1 ,1 ,1], [0.7, 1 ,1 ,1], 206 | [1, 0.05, 1, 1], [1, 0.1, 1, 1], [1, 0.15, 1, 1], [1, 0.2, 1, 1], [1, 0.3, 1, 1], 207 | [1, 1, 0.05, 1], [1, 1, 0.1, 1], [1, 1, 0.15, 1], [1, 1, 0.2, 1], [1, 1, 0.3, 1], 208 | [1, 1, 1, 0.05], [1, 1, 1, 0.1], [1, 1, 1, 0.15], [1, 1, 1, 0.2], [1, 1, 1, 0.3]] 209 | 210 | # 验证修剪率 211 | #eval_prune_threshold(test_threshold_list, test_prototxt, caffemodel, prune_layers) 212 | 213 | threshold = [0.3, 0.1, 0.01, 0.2] 214 | prune(threshold, solver.test_nets[0], prune_layers) 215 | pruned_model = model_dir + model_name +'_pruned' + '.caffemodel' 216 | solver.test_nets[0].save(pruned_model) 217 | 218 | retrain_pruned(solver, pruned_model, threshold, prune_layers) 219 | 220 | 221 | 222 | """ 223 | # 各层对应的修剪率 224 | threshold = [0.3, 0.1, 0.01, 0.2] 225 | 226 | net = caffe.Net(deploy, caffemodel, caffe.TEST) 227 | # 修剪 228 | prune(threshold, net, prune_layers, test_prototxt) 229 | 230 | # 保存修剪后的稀疏网络模型 231 | output_model = model_name +'_pruned' + '.caffemodel' 232 | net.save(output_model) 233 | """ -------------------------------------------------------------------------------- /quantize.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 通过Kmeans聚类的方法来量化权重 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | import scipy.cluster.vq as scv 6 | import pickle 7 | import os 8 | os.environ['GLOG_minloglevel'] = '2' 9 | import caffe 10 | import time 11 | 12 | # 获得各层的量化码表 13 | def kmeans_net(net, layers, num_c=16, initials=None): 14 | # net: 网络 15 | # layers: 需要量化的层 16 | # num_c: 各层的量化级别 17 | # initials: 初始聚类中心 18 | codebook = {} # 量化码表 19 | if type(num_c) == type(1): 20 | num_c = [num_c] * len(layers) 21 | else: 22 | assert len(num_c) == len(layers) 23 | 24 | # 对各层进行聚类分析 25 | print "==============Perform K-means=============" 26 | for idx, layer in enumerate(layers): 27 | print "Eval layer:", layer 28 | W = net.params[layer][0].data.flatten() 29 | W = W[np.where(W != 0)] # 筛选不为0的权重 30 | # 默认情况下，聚类中心为线性分布中心 31 | if initials is None: # Default: uniform sample 32 | min_W = np.min(W) 33 | max_W = np.max(W) 34 | initial_uni = np.linspace(min_W, max_W, num_c[idx] - 1) 35 | codebook[layer], _ = scv.kmeans(W, initial_uni) 36 | elif type(initials) == type(np.array([])): 37 | codebook[layer], _ = scv.kmeans(W, initials) 38 | elif initials == 'random': 39 | codebook[layer], _ = scv.kmeans(W, num_c[idx] - 1) 40 | else: 41 | raise Exception 42 | 43 | # 将0权重值附上 44 | codebook[layer] = np.append(0.0, codebook[layer]) 45 | print "codebook size:", len(codebook[layer]) 46 | 47 | return codebook 48 | 49 | # 随机量化权重值 50 | def stochasitc_quantize2(W, codebook): 51 | # mask插入新维度：(W.shape,1) 52 | mask = W[:, np.newaxis] - codebook 53 | 54 | mask_neg = mask 55 | mask_neg[mask_neg > 0.0] -= 99999.0 56 | max_neg = np.max(mask_neg, axis=1) 57 | max_code = np.argmax(mask_neg, axis=1) 58 | 59 | mask_pos = mask 60 | mask_pos += 99999.0 61 | min_code = np.argmin(mask_pos, axis=1) 62 | min_pos = np.min(mask_pos, axis=1) 63 | 64 | rd = np.random.uniform(low=0.0, high=1.0, size=(len(W))) 65 | thresh = min_pos.astype(np.float32) / (min_pos - max_neg) 66 | 67 | max_idx = thresh < rd 68 | min_idx = thresh >= rd 69 | 70 | codes = np.zeros(W.shape) 71 | codes[max_idx] += min_code[max_idx] 72 | codes[min_idx] += max_code[min_idx] 73 | 74 | return codes.astype(np.int) 75 | 76 | # 得到网络的量化权重值 77 | def quantize_net(net, codebook): 78 | layers = codebook.keys() 79 | codes_W = {} 80 | print "================Perform quantization==============" 81 | for layer in layers: 82 | print "Quantize layer:", layer 83 | W = net.params[layer][0].data 84 | codes, _ = scv.vq(W.flatten(), codebook[layer]) # 根据码表得到量化权重值 85 | # codes = stochasitc_quantize2(W.flatten(), codebook[layer]) # 采用随机量化的方式 86 | codes = np.reshape(codes, W.shape) 87 | codes_W[layer] = np.array(codes, dtype=np.uint32) 88 | # 将量化后的权重保存到网络中 89 | W_q = np.reshape(codebook[layer][codes], W.shape) 90 | np.copyto(net.params[layer][0].data, W_q) 91 | 92 | return codes_W 93 | 94 | 95 | def quantize_net_with_dict(net, layers, codebook, use_stochastic=False, timing=False): 96 | start_time = time.time() 97 | codeDict = {} # 记录各个量化中心所处的位置 98 | maskCode = {} # 各层量化结果 99 | for layer in layers: 100 | print "Quantize layer:", layer 101 | W = net.params[layer][0].data 102 | if use_stochastic: 103 | codes = stochasitc_quantize2(W.flatten(), codebook[layer]) 104 | else: 105 | codes, _ = scv.vq(W.flatten(), codebook[layer]) 106 | W_q = np.reshape(codebook[layer][codes], W.shape) 107 | net.params[layer][0].data[...] = W_q 108 | 109 | maskCode[layer] = np.reshape(codes, W.shape) 110 | codeBookSize = len(codebook[layer]) 111 | a = maskCode[layer].flatten() 112 | b = xrange(len(a)) 113 | 114 | codeDict[layer] = {} 115 | for i in xrange(len(a)): 116 | codeDict[layer].setdefault(a[i], []).append(b[i]) 117 | 118 | if timing: 119 | print "Update codebook time:%f" % (time.time() - start_time) 120 | 121 | return codeDict, maskCode 122 | 123 | def static_vars(**kwargs): 124 | def decorate(func): 125 | for k in kwargs: 126 | setattr(func, k, kwargs[k]) 127 | return func 128 | return decorate 129 | 130 | @static_vars(step_cache={}, step_cache2={}, count=0) 131 | def update_codebook_net(net, codebook, codeDict, maskCode, args, update_layers=None, snapshot=None): 132 | 133 | start_time = time.time() 134 | extra_lr = args['lr'] # 基础学习速率 135 | decay_rate = args['decay_rate'] # 衰减速率 136 | momentum = args['momentum'] # 遗忘因子 137 | update_method = args['update'] # 更新方法 138 | smooth_eps = 0 139 | 140 | normalize_flag = args['normalize_flag'] # 是否进行归一化 141 | 142 | 143 | if update_method == 'rmsprop': 144 | extra_lr /= 100 145 | 146 | # 对码表与量化结果的初始化 147 | if update_codebook_net.count == 0: 148 | step_cache2 = update_codebook_net.step_cache2 149 | step_cache = update_codebook_net.step_cache 150 | if update_method == 'adadelta': 151 | for layer in update_layers: 152 | step_cache2[layer] = {} 153 | for code in xrange(1, len(codebook[layer])): 154 | step_cache2[layer][code] = 0.0 155 | smooth_eps = 1e-8 156 | 157 | for layer in update_layers: 158 | step_cache[layer] = {} 159 | for code in xrange(1, len(codebook[layer])): 160 | step_cache[layer][code] = 0.0 161 | 162 | update_codebook_net.count = 1 163 | 164 | else: 165 | # 读入上次运算的结果 166 | step_cache2 = update_codebook_net.step_cache2 167 | step_cache = update_codebook_net.step_cache 168 | update_codebook_net.count += 1 169 | 170 | # 所有层名 171 | total_layers = net.params.keys() 172 | if update_layers is None: # 所有层都需要进行更新 173 | update_layers = total_layers 174 | 175 | # 权重码表的更新 176 | for layer in total_layers: 177 | if layer in update_layers: 178 | diff = net.params[layer][0].diff.flatten() # 误差梯度 179 | codeBookSize = len(codebook[layer]) 180 | dx = np.zeros((codeBookSize)) # 编码表的误差更新 181 | for code in xrange(1, codeBookSize): 182 | indexes = codeDict[layer][code] # codeDict保存属于某编码的权重的序号 183 | #diff_ave = np.sum(diff[indexes]) / len(indexes) 184 | diff_ave = np.sum(diff[indexes]) # 统计该编码所有的误差更新和 185 | 186 | # 针对于不同方法进行更新 187 | if update_method == 'sgd': 188 | dx[code] = -extra_lr * diff_ave 189 | elif update_method == 'momentum': 190 | if code in step_cache[layer]: 191 | dx[code] = momentum * step_cache[layer][code] - (1 - momentum) * extra_lr * diff_ave 192 | step_cache[layer][code] = dx 193 | elif update_method == 'rmsprop': 194 | if code in step_cache[layer]: 195 | step_cache[layer][code] = decay_rate * step_cache[layer][code] + (1.0 - decay_rate) * diff_ave ** 2 196 | dx[code] = -(extra_lr * diff_ave) / np.sqrt(step_cache[layer][code] + 1e-6) 197 | elif update_method == 'adadelta': 198 | if code in step_cache[layer]: 199 | step_cache[layer][code] = step_cache[layer][code] * decay_rate + (1.0 - decay_rate) * diff_ave ** 2 200 | dx[code] = -np.sqrt((step_cache2[layer][code] + smooth_eps) / (step_cache[layer][code] + smooth_eps)) * diff_ave 201 | step_cache2[layer][code] = step_cache2[layer][code] * decay_rate + (1.0 - decay_rate) * (dx[code] ** 2) 202 | 203 | # 是否需要进行归一化更新参数 204 | if normalize_flag: 205 | codebook[layer] += extra_lr * np.sqrt(np.mean(codebook[layer] ** 2)) / np.sqrt(np.mean(dx ** 2)) * dx 206 | else: 207 | codebook[layer] += dx 208 | else: 209 | pass 210 | 211 | # maskCode保存编码结果 212 | W2 = codebook[layer][maskCode[layer]] 213 | net.params[layer][0].data[...] = W2 # 量化后权重值 214 | 215 | print "Update codebook time:%f" % (time.time() - start_time) 216 | 217 | # 保存量化结果 218 | def store_all(net, codebook, dir_t, idx=0): 219 | net.save(dir_t + 'caffemodel%d' % idx) 220 | # 量化网络及码表 221 | pickle.dump(codebook, open(dir_t + 'codebook%d' % idx, 'w')) 222 | 223 | # 恢复权重值 224 | def recover_all(net, dir_t, idx=0): 225 | layers = net.params.keys() 226 | net.copy_from(dir_t + 'caffemodel%d' % idx) 227 | codebook = pickle.load(open(dir_t + 'codebook%d' % idx)) 228 | maskCode = {} 229 | codeDict = {} 230 | for layer in layers: 231 | W = net.params[layer][0].data 232 | # 码表结果 233 | codes, _ = scv.vq(W.flatten(), codebook[layer]) 234 | # 编码结果重新排列 235 | maskCode[layer] = np.reshape(codes, W.shape) 236 | codeBookSize = len(codebook[layer]) 237 | a = maskCode[layer].flatten() 238 | b = xrange(len(a)) 239 | 240 | codeDict[layer] = {} 241 | for i in xrange(len(a)): 242 | # codeDict保存每个码有哪些位置，而maskCode保存每个位置属于哪个码 243 | codeDict[layer].setdefault(a[i], []).append(b[i]) 244 | 245 | return codebook, maskCode, codeDict 246 | 247 | 248 | def analyze_log(fileName): 249 | data = open(fileName, "r") 250 | y = [] 251 | for line in data: 252 | y.append(float(line.split()[0])) 253 | return y 254 | 255 | # 读入测试数据 256 | def parse_caffe_log(log): 257 | lines = open(log).readlines() 258 | try: 259 | res = map(lambda x: float(x.split()[-1]), lines[-3:-1]) 260 | except Exception as e: 261 | print e 262 | res = [0.0, 0.0] 263 | return res 264 | 265 | # 检测量化后网络的精度 266 | def test_quantize_accu(test_net): 267 | test_iter = 100 268 | test_loss = 0 269 | accuracy = 0 270 | for test_it in range(test_iter): 271 | # 进行一次测试 272 | test_net.forward() 273 | # 计算test loss 274 | test_loss += test_net.blobs['loss'].data 275 | # 计算test accuracy 276 | accuracy += test_net.blobs['accuracy'].data 277 | 278 | return (test_loss / test_iter), (accuracy / test_iter) 279 | 280 | 281 | def save_quantize_net(codebook, maskcode, net_filename, total_layers): 282 | # 编码 283 | quantizeNet = {} 284 | for layer in total_layers: 285 | quantizeNet[layer+'_codebook'] = np.float32(codebook[layer]) 286 | quantizeNet[layer + '_maskcode'] = np.int8(maskcode[layer]) 287 | 288 | np.savez(net_filename,quantizeNet) 289 | 290 | # 保存修剪量化的网络参数 291 | def save_pruned_quantize_net(codebook, maskcode, net_filename, total_layers): 292 | # W_flatten: 扁平化的权重矩阵 293 | # num_level: 量化级别 294 | quantizeNet = {} 295 | for layer in total_layers: 296 | W_flatten = maskCode[layer].flatten() 297 | indx = 0 298 | num_level = 8 299 | csc_W = [] 300 | csc_indx = [] 301 | for n in range(len(W_flatten)): 302 | if W_flatten[n]!=0 or indx == 2**num_level: 303 | csc_W.append(W_flatten[n]) 304 | csc_indx.append(indx) 305 | indx = 0 306 | else: 307 | indx += 1 308 | if indx!=0: 309 | csc_W.append(0) 310 | csc_indx.append(indx-1) 311 | print max(csc_indx) 312 | quantizeNet[layer + '_codebook'] = np.float32(codebook[layer]) 313 | quantizeNet[layer + '_maskcode_W'] = np.array(csc_W, dtype=np.int8) 314 | print max(csc_indx) 315 | quantizeNet[layer + '_maskcode_indx'] = np.array(csc_indx, dtype=np.int8) 316 | 317 | np.savez(net_filename, quantizeNet) 318 | 319 | 320 | caffe.set_mode_gpu() 321 | caffe.set_device(0) 322 | 323 | caffe_root = '../../' 324 | model_dir = caffe_root + 'models/mnist/' 325 | deploy = model_dir + 'deploy.prototxt' 326 | solver_file = model_dir + 'solver.prototxt' 327 | # model_name = 'LeNet5_Mnist_shapshot_iter_10000' 328 | model_name = 'LeNet5_Mnist_shapshot_iter_10000_pruned' 329 | caffemodel = model_dir + model_name + '.caffemodel' 330 | 331 | dir_t = '/weight_quantize/' 332 | 333 | # 运行测试命令 334 | args = dict(lr=0.01, decay_rate = 0.0009, momentum = 0.9, update = 'adadelta', normalize_flag = False) 335 | 336 | start_time = time.time() 337 | 338 | solver = caffe.SGDSolver(solver_file) 339 | solver.net.copy_from(caffemodel) 340 | # 需要量化的权重 341 | total_layers = ['conv1','conv2','ip1','ip2'] 342 | 343 | num_c = 2 ** 8 # 量化级别，由8位整数表示 344 | codebook = kmeans_net(solver.test_nets[0], total_layers, num_c) 345 | 346 | codeDict, maskCode = quantize_net_with_dict(solver.test_nets[0], total_layers, codebook) 347 | quantize_net_caffemodel = model_dir + model_name + '_quantize.caffemodel' 348 | solver.test_nets[0].save(quantize_net_caffemodel) 349 | 350 | quantize_net_npz = model_dir + model_name + '_quantize_net' 351 | save_pruned_quantize_net(codebook, maskCode, quantize_net_npz , total_layers) 352 | 353 | # 迭代训练编码表 354 | accuracys = [] 355 | co_iters = 40 356 | ac_iters = 10 357 | for i in xrange(2500): 358 | if (i % (co_iters + ac_iters) == 0 and i > 0): 359 | # 重新量化 360 | # 导入训练后的 361 | codebook = kmeans_net(solver.net, total_layers, num_c) 362 | codeDict, maskCode = quantize_net_with_dict(solver.net, total_layers, codebook) 363 | solver.net.save(quantize_net_caffemodel) 364 | solver.test_nets[0].copy_from(quantize_net_caffemodel) 365 | _, accu = test_quantize_accu(solver.test_nets[0]) 366 | accuracys.append(accu) 367 | 368 | solver.step(1) 369 | if (i % (co_iters + ac_iters) < co_iters): 370 | # 码表更新 371 | update_codebook_net(solver.net, codebook, codeDict, maskCode, args=args, update_layers=total_layers) 372 | 373 | print "Iter:%d, Time cost:%f" % (i, time.time() - start_time) 374 | 375 | plt.plot(accuracys, 'r.-') 376 | plt.show() -------------------------------------------------------------------------------- /solver.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 生成solver文件 3 | from caffe.proto import caffe_pb2 4 | 5 | def solver_file(model_root, model_name): 6 | s = caffe_pb2.SolverParameter() # 声明solver结构 7 | s.train_net = model_root+'train.prototxt' # 训练网络结构配置文件 8 | s.test_net.append(model_root+'test.prototxt') # 测试时网络结构配置文件，测试网络可有多个 9 | # 每训练迭代test_interval次进行一次测试。 10 | s.test_interval = 500 11 | # 每次测试时的批量数，测试里网络可有多个 12 | s.test_iter.append(100) 13 | # 最大训练迭代次数 14 | s.max_iter = 10000 15 | # 基础学习率 16 | s.base_lr = 0.01 17 | # 动量，记忆因子 18 | s.momentum = 0.9 19 | # 权重衰减值，遗忘因子 20 | s.weight_decay = 5e-4 21 | # 学习率变化策略。可选参数：fixed、step、exp、inv、multistep 22 | # fixed: 保持base_lr不变； 23 | # step: 学习率变化规律base_lr * gamma ^ (floor(iter / stepsize))，其中iter表示当前的迭代次数； 24 | # exp: 学习率变化规律base_lr * gamma ^ iter； 25 | # inv: 还需要设置一个power，学习率变化规律base_lr * (1 + gamma * iter) ^ (- power)； 26 | # multistep: 还需要设置一个stepvalue，这个参数和step相似，step是均匀等间隔变化，而multistep则是根据stepvalue值变化； 27 | # stepvalue参数说明： 28 | # poly: 学习率进行多项式误差，返回base_lr (1 - iter/max_iter) ^ (power)； 29 | # sigmoid: 学习率进行sigmod衰减，返回base_lr ( 1/(1 + exp(-gamma * (iter - stepsize))))。 30 | s.lr_policy = 'inv' 31 | s.gamma = 0.0001 32 | s.power = 0.75 33 | 34 | s.display = 100 # 每迭代display次显示结果 35 | s.snapshot = 5000 # 保存临时模型的迭代数 36 | s.snapshot_prefix = model_root+model_name+'shapshot' # 模型前缀，就是训练好生成model的名字 37 | s.type = 'SGD' # 训练方法（各类梯度下降法），可选参数：SGD，AdaDelta，AdaGrad，Adam，Nesterov，RMSProp 38 | s.solver_mode = caffe_pb2.SolverParameter.GPU # 训练及测试模型，GPU或CPU 39 | 40 | solver_file=model_root+'solver.prototxt' # 要保存的solver文件名 41 | 42 | with open(solver_file, 'w') as f: 43 | f.write(str(s)) 44 | 45 | caffe_root = '../../' 46 | model_name = 'LeNet5_Mnist_' 47 | # caffe_root = 'E:/Code/Github/windows_caffe/' 48 | model_root = caffe_root + 'models/mnist/' 49 | solver_file(model_root, model_name) 50 | 51 | -------------------------------------------------------------------------------- /train_val.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # 训练及测试文件 3 | # 训练网络 4 | import caffe 5 | import numpy as np 6 | import matplotlib.pyplot as plt 7 | import math 8 | 9 | def crop_network(prune_proto, caffemodel, prune_caffemodel): 10 | # 截取已知网络的部分层 11 | # caffemodel网络权重值并不要求其结构与proto相对应 12 | # 网络只会取train_proto中定义的结构中权重作为网络的初始权重值 13 | # 因此，当我们需要截取某些已训练网络的特定层作为新网络的某些层的权重初始值，只需要在其train_proto定义同名的层 14 | # 之后caffe将在caffemodel中找到与train_proto定义的同名结构，并将其权重作为应用权重初始值。 15 | # prune_deploy: 选择保留的网络结构层:prototxt 16 | # caffemodel: 已知网络的权重连接 17 | # prune_caffemodel：截断网络的权重连接文件 18 | net = caffe.Net(prune_proto, caffemodel, caffe.TEST) 19 | net.save(prune_caffemodel) 20 | 21 | def train(solver_proto, caffemodel='', is_step=True, savefig=''): 22 | # 训练模型函数 23 | # solver_proto: 训练配置文件 24 | # caffemodel：预设权重值或者快照等，并不要求其结构与网络结构相对应，但只会取与训练网络结构相对应的权重值 25 | # is_step: True表示按步训练，False表示直接完成训练 26 | # savefig: 表示要保存的图像训练时损失变化图 27 | # 设置训练器：随机梯度下降算法 28 | solver = caffe.SGDSolver(solver_proto) 29 | if caffemodel!='': 30 | solver.net.copy_from(caffemodel) 31 | 32 | if is_step==False: 33 | # 直接完成训练 34 | solver.solve() 35 | else: 36 | # 迭代次数 37 | max_iter = 10000 38 | # 每隔100次收集一次数据 39 | display = 100 40 | 41 | # 每次测试进行100次解算，10000/100 42 | test_iter = 100 43 | # 每500次训练进行一次测试（100次解算），60000/64 44 | test_interval = 500 45 | 46 | # 初始化 47 | train_loss = np.zeros(int(math.ceil(max_iter * 1.0 / display))) 48 | test_loss = np.zeros(int(math.ceil(max_iter * 1.0 / test_interval))) 49 | test_acc = np.zeros(int(math.ceil(max_iter * 1.0 / test_interval))) 50 | 51 | # iteration 0，不计入 52 | solver.step(1) 53 | 54 | # 辅助变量 55 | _train_loss = 0 56 | _test_loss = 0 57 | _accuracy = 0 58 | 59 | # 分步训练 60 | for it in range(max_iter): 61 | # 进行一次解算 62 | solver.step(1) 63 | # 每迭代一次，训练batch_size张图片 64 | _train_loss += solver.net.blobs['loss'].data # 最后一层的损失值 65 | if it % display == 0: 66 | # 计算平均train loss 67 | train_loss[int(it / display)] = _train_loss / display 68 | _train_loss = 0 69 | 70 | # 测试 71 | if it % test_interval == 0: 72 | for test_it in range(test_iter): 73 | # 进行一次测试 74 | solver.test_nets[0].forward() 75 | # 计算test loss 76 | _test_loss += solver.test_nets[0].blobs['loss'].data 77 | # 计算test accuracy 78 | _accuracy += solver.test_nets[0].blobs['accuracy'].data 79 | # 计算平均test loss 80 | test_loss[it / test_interval] = _test_loss / test_iter 81 | # 计算平均test accuracy 82 | test_acc[it / test_interval] = _accuracy / test_iter 83 | _test_loss = 0 84 | _accuracy = 0 85 | 86 | # 绘制train loss、test loss和accuracy曲线 87 | print '\nplot the train loss and test accuracy\n' 88 | _, ax1 = plt.subplots() 89 | ax2 = ax1.twinx() 90 | 91 | # train loss -> 绿色 92 | ax1.plot(display * np.arange(len(train_loss)), train_loss, 'g') 93 | # test loss -> 黄色 94 | ax1.plot(test_interval * np.arange(len(test_loss)), test_loss, 'y') 95 | # test accuracy -> 红色 96 | ax2.plot(test_interval * np.arange(len(test_acc)), test_acc, 'r') 97 | 98 | ax1.set_xlabel('iteration') 99 | ax1.set_ylabel('loss') 100 | ax2.set_ylabel('accuracy') 101 | 102 | if savefig!='': 103 | plt.savefig(savefig) 104 | plt.show() 105 | 106 | #CPU或GPU模型转换 107 | #caffe.set_mode_cpu() 108 | caffe.set_device(0) 109 | caffe.set_mode_gpu() 110 | 111 | caffe_root = '../../' 112 | # caffe_root = 'E:/Code/Github/windows_caffe/' 113 | model_root = caffe_root + 'models/mnist/' 114 | solver_proto = model_root + 'solver.prototxt' 115 | train(solver_proto, caffemodel='', is_step=True) 116 | --------------------------------------------------------------------------------