pytorch imagenet测试代码
image_test.py
import argparse import numpy as np import sys import os import csv from imagenet_test_base import TestKit import torch class TestTorch(TestKit): def __init__(self): super(TestTorch, self).__init__() self.truth['tensorflow']['inception_v3'] = [(22, 9.6691055), (24, 4.3524747), (25, 3.5957973), (132, 3.5657473), (23, 3.346283)] self.truth['keras']['inception_v3'] = [(21, 0.93430489), (23, 0.002883445), (131, 0.0014781791), (24, 0.0014518998), (22, 0.0014435351)] self.model = self.MainModel.KitModel(self.args.w) self.model.eval() def preprocess(self, image_path): x = super(TestTorch, self).preprocess(image_path) x = np.transpose(x, (2, 0, 1)) x = np.expand_dims(x, 0).copy() self.data = torch.from_numpy(x) self.data = torch.autograd.Variable(self.data, requires_grad = False) def print_result(self, image_name, top1, top5): predict = self.model(self.data) predict = predict.data.numpy() return super(TestTorch, self).print_result(predict, image_name, top1, top5) def print_intermediate_result(self, layer_name, if_transpose=False): intermediate_output = self.model.test.data.numpy() super(TestTorch, self).print_intermediate_result(intermediate_output, if_transpose) def inference(self, images): with open(images) as fp_images: images_file = csv.reader(fp_images, delimiter='\n') top1 = 0.0 top5 = 0.0 image_count = 0 for image_name in images_file: image_count += 1 image_path = "../data/imagenet/small_imagenet/"+image_name[0] self.preprocess(image_path) temp1, temp5 = self.print_result(image_name[0], top1, top5) top1 = temp1 top5 = temp5 print("top1's accuracy : %f"%(top1/image_count)) print("top5's accuracy : %f"%(top5/image_count)) # self.print_intermediate_result(None, False) # self.test_truth() def dump(self, path=None): if path is None: path = self.args.dump torch.save(self.model, path) print('PyTorch model file is saved as [{}], generated by [{}.py] and [{}].'.format( path, self.args.n, self.args.w)) if __name__=='__main__': tester = TestTorch() if tester.args.dump: tester.dump() else: tester.inference(tester.args.image)
image_test_base.py:
请见上传的代码。 下载地址:https://files.cnblogs.com/files/jzcbest1016/imagenet_test_base.py.tar.gz
执行py文件时,需要终端输入参数:
parser = argparse.ArgumentParser() parser.add_argument('-p', '--preprocess', type=_text_type, help='Model Preprocess Type') # pytorch的测试程序, 这里为image_test.py parser.add_argument('-n', type=_text_type, default='kit_imagenet', help='Network structure file name.') # 模型结构测试文件 parser.add_argument('-s', type=_text_type, help='Source Framework Type', choices=self.truth.keys()) # 框架类型:pytorch,tensorflow... parser.add_argument('-w', type=_text_type, required=True, help='Network weights file name') #模型结构文件 parser.add_argument('--image', '-i', type=_text_type, help='Test image path.', default="../data/file_list.txt" #图像路径 ) parser.add_argument('-l', '--label', type=_text_type, default='../data/val.txt', help='Path of label.') #测试集类别 parser.add_argument('--dump', type=_text_type, default=None, help='Target model path.') # 转化的目标模型文件的保存路径 parser.add_argument('--detect', type=_text_type, default=None, help='Model detection result path.') # tensorflow dump tag parser.add_argument('--dump_tag', type=_text_type, default=None, help='Tensorflow model dump type', choices=['SERVING', 'TRAINING'])