MXNet ImageNet-11k

安装

先从github下载。

git clone --recursive https://github.com/dmlc/mxnet

接着可以选择二种方式来进行安装。

一种: make

cd mxnet
cp ./make/config.mk .
make

另一种: python

cd mxnet/python
python setup.py install

不过用第二种方法的话，还是无法import。

所以使用第一种方法，接着用下面语句import即可。

import sys
sys.path.insert(0, '这里写mxnet/python路径')
import mxnet as mx

下载训练好的模型

def download(url,prefix=''):
    filename = prefix+url.split("/")[-1]
    if not os.path.exists(filename):
        urllib.urlretrieve(url,filename)

path='http://data.mxnet.io/models/imagenet-11k/'
download(path+'resnet-152/resnet-152-symbol.json','full-')
download(path+'resnet-152/resnet-152-0000.params','full-')
download(path+'synset.txt','full-')

* 进行预测*

我们先对小猫进行预测。

with open('full-synset.txt','r') as file:
    synsets = [l.rstrip() for l in file]

# print synsets
sym, arg_params,aux_params = mx.model.load_checkpoint('full-resnet-152',0)

mod= mx.mod.Module(symbol=sym, context=mx.gpu())
mod.bind(for_training=False, data_shapes=[('data', (1,3,224,224))])
mod.set_params(arg_params, aux_params)

import cv2
import numpy as np
from collections import namedtuple
Batch= namedtuple('Batch', ['data'])

print Batch

def predict(filename, mod, synsets):
    img = cv2.imread(filename)
    if img is None:
        return None
    img = cv2.resize(img, (224,224))
    img = np.swapaxes(img,0,2)
    img = np.swapaxes(img,1,2)
    img = img[np.newaxis, :] 

    mod.forward(Batch([mx.nd.array(img)]))
    prob = mod.get_outputs()[0].asnumpy()
    prob = np.squeeze(prob)

    a = np.argsort(prob)[::-1]    
    for i in a[0:5]:
        print('probability=%f, class=%s'%(prob[i], synsets[i]))

predict('/home/lol/dl/image/cat/cat.jpg', mod, synsets)

接着随便从网上找个图。

predict('/home/lol/dl/image/strawberry/strawberry.jpg', mod, synsets)