pytorch模型调用和速度测试

硬件：

 

1、测试模型：NIMA的pytorch版本

来源：https://github.com/truskovskiyk/nima.pytorch

模型：pretrain-model.pth

速度：平均一张图300ms左右

5.152941832318902
5.508232474792749
4.2969538709148765
5.8256916594691575
5.338452965952456
5.439441880211234
4.638633335009217
avg time:  304.443359375  ms

测试调用代码：

from nima.inference.inference_model import InferenceModel
import os
import time

if __name__ == '__main__':
    path_model = '/Users/xuqiong/code/image/NIMA/nima.pytorch/pretrain-model.pth'
    model = InferenceModel(path_to_model=path_model)

    dir = '/Users/xuqiong/code/image/NIMA/nima.pytorch/test2'
    #dir = '/Users/xuqiong/code/image/data/tid2013/distorted_images'
    filelist = []
    filenames = os.listdir(dir)
    timeall = 0
    for fn in filenames:
        if fn == '.DS_Store':
            continue
        fullfilename = os.path.join(dir, fn)
        filelist.append(fullfilename)

    for i in range(0, len(filelist)):
        path_image = filelist[i]

        try:
         time0 = time.time()
         result = model.predict_from_file(path_image)#1-10分
         time1 = time.time()
         timed = time1 - time0
         timeall = timeall + timed
         print(result["mean_score"])
        except:
            continue
    print("avg time: ", timeall*1000/len(filelist), " ms")

 

2、测试模型：shufflenetv2

来源：https://github.com/miaow1988/ShuffleNet_V2_pytorch_caffe/releases

模型：shufflenet_v2_x0.25

速度：平均18ms

18.815040588378906
22.640228271484375
17.702817916870117
17.43292808532715
17.948150634765625
17.450809478759766
17.821073532104492
avg time:  18.544435501098633  ms

测试调用代码：

#!/usr/bin/env python
# encoding: utf-8

import torch
import torch.nn as nn
import os
import time

import slim
from slim import g_name

class BasicBlock(nn.Module):

    def __init__(self, name, in_channels, out_channels, stride, dilation):
        super(BasicBlock, self).__init__()
        self.g_name = name
        self.in_channels = in_channels
        self.stride = stride
        channels = out_channels//2
        if stride == 1:
            assert in_channels == out_channels
            self.conv = nn.Sequential(
                slim.conv_bn_relu(name + '/conv1', channels, channels, 1),
                slim.conv_bn(name + '/conv2',
                    channels, channels, 3, stride=stride,
                    dilation=dilation, padding=dilation, groups=channels),
                slim.conv_bn_relu(name + '/conv3', channels, channels, 1),
            )
        else:
            self.conv = nn.Sequential(
                slim.conv_bn_relu(name + '/conv1', in_channels, channels, 1),
                slim.conv_bn(name + '/conv2',
                    channels, channels, 3, stride=stride,
                    dilation=dilation, padding=dilation, groups=channels),
                slim.conv_bn_relu(name + '/conv3', channels, channels, 1),
            )
            self.conv0 = nn.Sequential(
                slim.conv_bn(name + '/conv4',
                    in_channels, in_channels, 3, stride=stride,
                    dilation=dilation, padding=dilation, groups=in_channels),
                slim.conv_bn_relu(name + '/conv5', in_channels, channels, 1),
            )
        self.shuffle = slim.channel_shuffle(name + '/shuffle', 2)

    def forward(self, x):
        if self.stride == 1:
            x1 = x[:, :(x.shape[1]//2), :, :]
            x2 = x[:, (x.shape[1]//2):, :, :]
            x = torch.cat((x1, self.conv(x2)), 1)
        else:
            x = torch.cat((self.conv0(x), self.conv(x)), 1)
        return self.shuffle(x)


class Network(nn.Module):

    def __init__(self, num_classes, width_multiplier):
        super(Network, self).__init__()
        assert width_multiplier in (0.25, 0.5, 1.0, 1.5, 2.0)
        self.num_classes = num_classes
        in_channels = 24
        width_config = {
            0.25: (24, 48, 96, 512),
            0.5: (48, 96, 192, 1024),
            1.0: (116, 232, 464, 1024),
            1.5: (176, 352, 704, 1024),
            2.0: (244, 488, 976, 2048),
        }

        # outputs, stride, dilation, blocks, type
        self.network_config = [
            g_name('data/bn', nn.BatchNorm2d(3)),
            slim.conv_bn_relu('stage1/conv', 3, in_channels, 3, 2, 1),
            # g_name('stage1/pool', nn.MaxPool2d(3, 2, 1)),
            g_name('stage1/pool', nn.MaxPool2d(3, 2, 0, ceil_mode=True)),
            (width_config[width_multiplier][0], 2, 1, 4, 'b'),
            (width_config[width_multiplier][1], 2, 1, 8, 'b'), # x16
            (width_config[width_multiplier][2], 2, 1, 4, 'b'), # x32
            slim.conv_bn_relu('conv5', width_config[width_multiplier][2], width_config[width_multiplier][3], 1),
            g_name('pool', nn.AvgPool2d(7, 1)),
            g_name('fc', nn.Conv2d(width_config[width_multiplier][3], self.num_classes, 1)),
        ]
        self.network = []
        for i, config in enumerate(self.network_config):
            if isinstance(config, nn.Module):
                self.network.append(config)
                continue
            out_channels, stride, dilation, num_blocks, stage_type = config
            stage_prefix = 'stage_{}'.format(i - 1)
            blocks = [BasicBlock(stage_prefix + '_1', in_channels,
                out_channels, stride, dilation)]
            for i in range(1, num_blocks):
                blocks.append(BasicBlock(stage_prefix + '_{}'.format(i + 1),
                    out_channels, out_channels, 1, dilation))
            self.network += [nn.Sequential(*blocks)]

            in_channels = out_channels
        self.network = nn.Sequential(*self.network)

        for name, m in self.named_modules():
            if any(map(lambda x: isinstance(m, x), [nn.Linear, nn.Conv1d, nn.Conv2d])):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def trainable_parameters(self):
        parameters = [
            {'params': self.cls_head_list.parameters(), 'lr_mult': 1.0},
            {'params': self.loc_head_list.parameters(), 'lr_mult': 1.0},
            # {'params': self.network.parameters(), 'lr_mult': 0.1},
        ]
        for i in range(len(self.network)):
            lr_mult = 0.1 if i in (0, 1, 2, 3, 4, 5) else 1
            parameters.append(
                {'params': self.network[i].parameters(), 'lr_mult': lr_mult}
            )
        return parameters

    def forward(self, x):
        x = self.network(x)
        return x.reshape(x.shape[0], -1)

class Model:
    def __init__(self, path_to_model):
        self.transform = Transform().val_transform
        self.model = NIMA(pretrained_base_model=False)
        state_dict = torch.load(path_to_model, map_location=lambda storage, loc: storage)
        self.model.load_state_dict(state_dict)
        self.model = self.model.to(device)
        self.model.eval()

    def predict_from_file(self, image_path):
        image = default_loader(image_path)
        return self.predict(image)

    def predict_from_pil_image(self, image):
        image = image.convert('RGB')
        return self.predict(image)

    def predict(self, image):
        image = self.transform(image)
        image = image.unsqueeze_(0)
        image = image.to(device)
        image = torch.autograd.Variable(image, volatile=True)
        prob = self.model(image).data.cpu().numpy()[0]

        mean_score = get_mean_score(prob)
        std_score = get_std_score(prob)

        return format_output(mean_score, std_score, prob)

if __name__ == '__main__':

    import PIL.Image
    import torchvision
    import numpy as np

    def assert_diff(a, b):
        if isinstance(a, torch.Tensor):
            a = a.detach().cpu().numpy()
        if isinstance(b, torch.Tensor):
            b = b.detach().cpu().numpy()
        print(a.shape, b.shape)
        a = a.reshape(-1)
        b = b.reshape(-1)
        assert a.shape == b.shape
        diff = np.abs(a - b)
        print('mean diff = %f' % diff.mean())
        assert diff.mean() < 0.001
        print('max diff = %f' % diff.max())
        assert diff.max() < 0.001

    # Initilize a PyTorch model.
    num_classes = 1000
    model_width = 0.25
    net = Network(num_classes, model_width).train(False)
    print(net)
    load_pytorch = '/Users/xuqiong/code/image/Pytorch/shufflenet_v2_x0.25.pth'
    net.load_state_dict(torch.load(load_pytorch, map_location=lambda storage, loc: storage))

    ##############test##############
    dir = '/Users/xuqiong/code/image/NIMA/nima.pytorch/test2'
    filelist = []
    filenames = os.listdir(dir)
    timeall = 0
    for fn in filenames:
        if fn == '.DS_Store':
            continue
        fullfilename = os.path.join(dir, fn)
        filelist.append(fullfilename)

    image_hw = 224
    image_hw = 224
    timeall = 0
    for i in range(0, len(filelist)):
        test = filelist[i]

        img = PIL.Image.open(test).convert('RGB')
        img = torchvision.transforms.functional.resize(img, (image_hw, image_hw))
        img = torchvision.transforms.functional.to_tensor(img).unsqueeze(0).numpy()
        x = torch.tensor(img.copy(), dtype=torch.float32)
        with torch.no_grad():
            time0 = time.time()
            cls_results = net(x)
            time1 = time.time()
            timed = time1 - time0
            timeall = timeall + timed
            print(timed*1000)
        #print("result: ", cls_results.shape)

    print("avg time: ", timeall*1000/len(filelist), " ms")