__init__.py 709 Bytes
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
"""Store for caffe2 examples and common models."""
from __future__ import absolute_import as _abs
import os
import sys
import importlib
from . import squeezenet
from caffe2.python.models.download import ModelDownloader

models = [
    'squeezenet',
    'resnet50',
    'vgg19',
]

mf = ModelDownloader()

class Model:
    def __init__(self, model_name):
        self.init_net, self.predict_net, self.value_info = mf.get_c2_model(model_name)

for model in models:
    try:
        locals()['c2_' + model] = importlib.import_module('caffe2.python.models.' + model)
    except ImportError:
        locals()['c2_' + model] = Model(model)

# squeezenet
def relay_squeezenet():
    return squeezenet.get_workload()