mxnet的Predict with pre-trained models教程，以及一些常见的问题

最新推荐文章于 2024-12-04 04:27:13 发布

xiaotao_1

最新推荐文章于 2024-12-04 04:27:13 发布

阅读量1.6k

点赞数

分类专栏： mxnet 文章标签： mxnet 预训练模型

本文链接：https://blog.youkuaiyun.com/xiaotao_1/article/details/78836206

版权

mxnet 专栏收录该内容

26 篇文章

订阅专栏

　　问题1：要下载的模型比较大，可以手动去http://data.mxnet.io/models/imagenet-11k/下载，然后把resnet-152-0000.params； resnet-152-symbol.json； synset.txt放到该.py文件所在的文件夹。
　　问题2：import cv2 .但是有的人没有装这个，参考http://blog.youkuaiyun.com/xiaotao_1/article/details/78834039解决。

　　以下是代码：

import cv2
import mxnet as mx
import matplotlib.pyplot as plt
import numpy as np
from collections import namedtuple


# define a simple data batch
Batch = namedtuple('Batch', ['data'])
sym, arg_params, aux_params = mx.model.load_checkpoint('resnet-152', 0)
mod = mx.mod.Module(symbol=sym, context=mx.gpu(), label_names=None)
mod.bind(for_training=False, data_shapes=[('data', (1,3,224,224))],
         label_shapes=mod._label_shapes)
mod.set_params(arg_params, aux_params, allow_missing=True)
with open('synset.txt', 'r') as f:
    labels = [l.rstrip() for l in f]

def get_image(url, show=False):
    # download and show the image
    fname = mx.test_utils.download(url)
    img = cv2.cvtColor(cv2.imread(fname), cv2.COLOR_BGR2RGB)
    if img is None:
         return None
    if show:
         plt.imshow(img)
         plt.axis('off')
         plt.show()
    # convert into format (batch, RGB, width, height)
    img = cv2.resize(img, (224, 224))
    img = np.swapaxes(img, 0, 2)
    img = np.swapaxes(img, 1, 2)
    img = img[np.newaxis, :]
    return img

def predict(url):
    img = get_image(url, show=True)
    # compute the predict probabilities
    mod.forward(Batch([mx.nd.array(img)]))
    prob = mod.get_outputs()[0].asnumpy()
    # print the top-5
    prob = np.squeeze(prob)
    a = np.argsort(prob)[::-1]
    for i in a[0:5]:
        print('probability=%f, class=%s' %(prob[i], labels[i]))
predict('http://writm.com/wp-content/uploads/2016/08/Cat-hd-wallpapers.jpg')

　　结果是：

probability=0.695960, class=n02122948 kitten, kitty
probability=0.042496, class=n01323155 kit
probability=0.030121, class=n01318894 pet
probability=0.029757, class=n02122878 tabby, queen
probability=0.026040, class=n01322221 baby