当我们训练完一个 caffe 的模型想要简单快速的测试一下模型在 CPU/GPU 下的图片检测速度时,可以用到下面的程序,此处就是把同一张图片循环 Forward 1000 次的方法来查看模型速度。
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import time
# caffe的根目录
caffe_root = '/home/t702/caffe/'
sys.path.insert(0,caffe_root + 'python')
import caffe
# deploy文件的路径
Model_FILE = '/PATH/TO/deploy.prototxt'
# 预训练好的caffemodel的模型
PRETRAINED = '/PATH/TO/XXXX.caffemodel'
# 测试图片的路径
IMAGE_FILE = '/PATH/TO/11.png'
input_image = caffe.io.load_image(IMAGE_FILE)
net = caffe.Classifier(Model_FILE,PRETRAINED)
#模式设置为CPU
caffe.set_mode_cpu()
#模式设置为GPU
#caffe.set_mode_gpu()
start = time.clock()
# 单张图片循环一千次测试时间
for i in range(0,1000):
prediction = net.predict([input_image],oversample=False)
end = time.clock()
print('predicted calss: ',prediction[0].argmax())
print(end-start)