test.py

import os
import sys
import imp
import numpy as np
import warnings
warnings.filterwarnings("ignore")

# Test for Torch
def torch(test_models, model_path, img_path):
    results_o, results_d, op_sets = dict(), dict(), dict()

    from PIL import Image
    import torch
    import torchvision.models as models
    from torchvision import transforms
    from torch.autograd import Variable

    # Torch to IR
    from ox.pytorch.pytorch_parser import PytorchParser

    for model in test_models:
        if 'inception' in model: image_size = 299
        else: image_size = 224

        image = Image.open(img_path)
        transformation = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        image_tensor = transformation(image).float()
        image_tensor = image_tensor.unsqueeze_(0)
        x = Variable(image_tensor)
        inputshape = [3, image_size, image_size]

        arch_filename = os.path.join(model_path, 'PyTorch', model+'.pth')

        # test model
        if 'resnet50' in model:
            model_eval = models.resnet50()
        elif 'inception' in model:
            from models.torch import inception
            model_eval = inception.inceptionresnetv2(pretrained=False)
        elif 'shufflenet' in model:
            from models.torch import shufflenet
            model_eval = shufflenet.shufflenet()
        elif 'fcn' in model:
            from models.torch import fcn
            model_eval = fcn.FCNs()
        elif 'lstm' in model:
            from models.torch import lstm
            model_eval = lstm.Lstm()

        model_eval.eval()
        predict = model_eval(x).data.numpy()
        preds = np.squeeze(predict)

        print('\033[1;31;40m')
        print(' Result of', model, ': ', np.argmax(preds))
        print('\033[0m')
        results_o[model] = preds
        torch.save(model_eval, arch_filename)

        # convert
        IR_filename = os.path.join(model_path, 'IR', model+'_torch')
        parser = PytorchParser(arch_filename, inputshape)
        ops = parser.run(IR_filename)
        op_sets[model] = ops
    del parser
    del PytorchParser

    # IR to Torch
    from ox.pytorch.pytorch_emitter import PytorchEmitter

    for model in test_models:
        if 'inception' in model: image_size = 299
        else: image_size = 224

        image = Image.open(img_path)
        transformation = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        image_tensor = transformation(image).float()
        image_tensor = image_tensor.unsqueeze_(0)
        x = Variable(image_tensor)
        inputshape = [3, image_size, image_size]

        arch_filename = os.path.join(model_path, 'IR', model+'_torch.pb')
        weight_filename = os.path.join(model_path, 'IR', model+'_torch.npy')
        converted_file = os.path.join(model_path, 'PyTorch', model+'_ox')
        emitter = PytorchEmitter((arch_filename, weight_filename))
        emitter.run(converted_file + '.py', converted_file + '.npy', 'test')

        model_converted = imp.load_source('PytorchModel', converted_file + '.py').KitModel(converted_file + '.npy')

        model_converted.eval()

        predict = model_converted(x).data.numpy()
        preds = np.squeeze(predict)

        print('\033[1;31;40m')
        print(' Result of ', model+'_ox : ', np.argmax(preds))
        print('\033[0m')
        results_d[model] = np.mean(np.power(results_o[model] - preds, 2))
    del emitter
    del PytorchEmitter

    return results_d, op_sets


# Test for Tensorflow
def tensorflow(test_models, model_path, img_path):
    results_o, results_d, op_sets = dict(), dict(), dict()

    import tensorflow as tf
    from PIL import Image
    image = Image.open(img_path)

    # Tensorflow to IR
    from ox.tensorflow.tensorflow_parser import TensorflowParser

    for model in test_models:
        arch_filename = os.path.join(model_path, 'tensorflow', model, model+'.ckpt.meta')
        weight_filename = os.path.join(model_path, 'tensorflow', model, model+'.ckpt')
        
        # test model
        if 'resnet50' in model:
            img = np.array(image.resize((299, 299), Image.ANTIALIAS))
            x = np.expand_dims(img, axis=0)
            from models.tf import resnet50
            preds = resnet50.test(x, model_path)
        elif 'inception' in model:
            img = np.array(image.resize((224, 224), Image.ANTIALIAS))
            x = np.expand_dims(img, axis=0)
            from models.tf import inception_v3
            preds = inception_v3.test(x, model_path)
        elif 'shufflenet' in model:
            img = np.array(image.resize((224, 224), Image.ANTIALIAS))
            x = np.expand_dims(img, axis=0)
            from models.tf import shufflenet
            preds = shufflenet.test(x, model_path)
        elif 'fcn' in model:
            img = np.array(image.resize((224, 224), Image.ANTIALIAS))
            x = np.expand_dims(img, axis=0)
            from models.tf import fcn
            preds = fcn.test(x, model_path)
        elif 'lstm' in model:
            img = np.array(image.resize((224, 224), Image.ANTIALIAS))
            x = np.expand_dims(img, axis=0)
            from models.tf import lstm
            preds = lstm.test(x, model_path)

        print('\033[1;31;40m')
        print(' Result of', model, ': ', np.argmax(preds))
        print('\033[0m')
        preds = np.squeeze(preds)
        if 'fcn' in model: preds = np.array(preds).astype(np.int32)
        results_o[model] = preds

        import tensorflow.contrib.keras as keras
        keras.backend.clear_session()

        # parser
        IR_filename = os.path.join(model_path, 'IR', model+'_tf')
        parser = TensorflowParser(arch_filename, weight_filename, ["OX_output"])
        ops = parser.run(IR_filename)
        op_sets[model] = ops
    del parser
    del TensorflowParser

    # IR to Tensorflow
    from ox.tensorflow.tensorflow_emitter import TensorflowEmitter

    for model in test_models:
        arch_filename = os.path.join(model_path, 'IR', model+'_tf.pb')
        weight_filename = os.path.join(model_path, 'IR', model+'_tf.npy')
        converted_file = os.path.join(model_path, 'tensorflow', model, model+'_ox')

        emitter = TensorflowEmitter((arch_filename, weight_filename))
        emitter.run(converted_file + '.py', None, 'test')

        # test model
        if 'resnet' in model:
            img = image.resize((299, 299), Image.ANTIALIAS)
        else:
            img = image.resize((224, 224), Image.ANTIALIAS)
        img = np.array(img)
        x = np.expand_dims(img, axis=0)
        if 'lstm' in model:
            x = np.reshape(x, (-1, 224 * 224 * 3))

        model_converted = imp.load_source('TFModel', converted_file + '.py').KitModel(weight_filename)

        input_tf, model_tf = model_converted

        with tf.Session() as sess:
            init = tf.global_variables_initializer()
            sess.run(init)
            predict = sess.run(model_tf, feed_dict = {input_tf : x})
        del model_converted
        del sys.modules['TFModel']
        preds = np.squeeze(predict)
        if 'fcn' in model: preds = np.array(preds).astype(np.int32)

        print('\033[1;31;40m')
        print(' Result of ', model+'_ox : ', np.argmax(preds))
        print('\033[0m')
        results_d[model] = np.mean(np.power(results_o[model] - preds, 2))

    del emitter
    del TensorflowEmitter

    return results_d, op_sets

def mk_dirs(path):
    if not os.path.exists(path):
        os.makedirs(path)
        return True
    return False

if __name__=='__main__':
    test_models = ['resnet50', 'inception_v3', 'shufflenet', 'fcn', 'lstm']
    model_path = './../models'
    img_path = os.path.join('./', 'elephant.jpg')

    # mkdirs
    mk_dirs(os.path.join(model_path, 'IR'))
    mk_dirs(os.path.join(model_path, 'tensorflow'))
    mk_dirs(os.path.join(model_path, 'PyTorch'))

    tf_err, tf_op_sets = tensorflow(test_models, model_path, img_path)
    torch_err, torch_op_sets = torch(test_models, model_path, img_path)

    for model in test_models:
        print('Model: {}'.format(model))
        print('- Error: tf ({}) | torch ({})'.format(tf_err[model], torch_err[model]))
        print('- TF Ops: {}'.format(tf_op_sets[model]))
        print('- Torch Ops: {}'.format(torch_op_sets[model]))
        print('\n')