predict.py

import numpy as np
import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), '../../build/python'))

from singa import device
from singa import tensor
import net as ffnet

def predict(net, images, cuda, topk=5):
    x = tensor.from_numpy(images.astype(np.float32))
    x.to_device(cuda)
    y = net.predict(x)
    y.to_host()
    y = tensor.to_numpy(y)
    prob = np.average(y, 0)
    labels = np.flipud(np.argsort(prob)) # sort prob in descending order
    return labels[0:topk], prob[labels[0:topk]]


def load_dataset(filepath):
    print 'Loading data file %s' % filepath
    with open(filepath, 'rb') as fd:
        cifar10 = cPickle.load(fd)
    image = cifar10['data'].astype(dtype=np.uint8)
    image = image.reshape((-1, 3, 32, 32))
    label = np.asarray(cifar10['labels'], dtype=np.uint8)
    label = label.reshape(label.size, 1)
    return image, label


def load_train_data(dir_path, num_batches=5):
    labels = []
    batchsize = 10000
    images = np.empty((num_batches * batchsize, 3, 32, 32), dtype=np.uint8)
    for did in range(1, num_batches + 1):
        fname_train_data = dir_path + "/data_batch_{}".format(did)
        image, label = load_dataset(fname_train_data)
        images[(did - 1) * batchsize:did * batchsize] = image
        labels.extend(label)
    images = np.array(images, dtype=np.float32)
    labels = np.array(labels, dtype=np.int32)
    return images, labels

def load_test_data(dir_path):
    images, labels = load_dataset(dir_path + "/test_batch")
    return np.array(images,  dtype=np.float32), np.array(labels, dtype=np.int32)


def compute_image_mean(train_dir):
    images = np.load(train_dir)
    return np.average(images, 0)

if __name__=='__main__':
    model = ffnet.create_alexnet()
    model.load('model.bin')
    cuda = device.create_cuda_gpu()
    model.to_device(cuda)

    mean = compute_image_mean('cifar-10-batches-py')
    test_images,_ = load_test_data('cifar-10-batches-py')
    print predict(model, test_images - mean, cuda)