Python caffe2.python.workspace.Predictor() Examples

def main():
    args = make_args()
    config = configparser.ConfigParser()
    utils.load_config(config, args.config)
    for cmd in args.modify:
        utils.modify_config(config, cmd)
    with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
        logging.config.dictConfig(yaml.load(f))
    torch.manual_seed(args.seed)
    model_dir = utils.get_model_dir(config)
    init_net = caffe2_pb2.NetDef()
    with open(os.path.join(model_dir, 'init_net.pb'), 'rb') as f:
        init_net.ParseFromString(f.read())
    predict_net = caffe2_pb2.NetDef()
    with open(os.path.join(model_dir, 'predict_net.pb'), 'rb') as f:
        predict_net.ParseFromString(f.read())
    p = workspace.Predictor(init_net, predict_net)
    height, width = tuple(map(int, config.get('image', 'size').split()))
    tensor = torch.randn(1, 3, height, width)
    # Checksum
    output = p.run([tensor.numpy()])
    for key, a in [
        ('tensor', tensor.cpu().numpy()),
        ('output', output[0]),
    ]:
        print('\t'.join(map(str, [key, a.shape, utils.abs_mean(a), hashlib.md5(a.tostring()).hexdigest()]))) 

def load_model(init_net_path, predict_net_path):
    with open(init_net_path, "rb") as f:
        init_net = f.read()
    with open(predict_net_path, "rb") as f:
        predict_net = f.read()
    p = workspace.Predictor(init_net, predict_net)
    return p 

def classify(self, path):
        input_image_size = self.model[4]

        img = skimage.img_as_float(skimage.io.imread(path)).astype(np.float32)
        img = self.rescale(img, input_image_size, input_image_size)
        img = self.crop_center(img, input_image_size, input_image_size)

        img = img.swapaxes(1, 2).swapaxes(0, 1)
        img = img[(2, 1, 0), :, :]
        img = img * 255 - self.mean

        img = img[np.newaxis, :, :, :].astype(np.float32)

        p = workspace.Predictor(self.init_net, self.predict_net)

        results = p.run([img])
        results = np.asarray(results)

        results = np.delete(results, 1)
        filtered_results = []

        for i, r in enumerate(results):
            if (float(r) > 0.01):
                filtered_results.append((self.get_category_from_code(i + 1), float(r)))

        return sorted(filtered_results, key=lambda result: result[1], reverse=True) 

def generate_test_output_data(caffe2_init_net, caffe2_predict_net, inputs):
    p = c2_workspace.Predictor(caffe2_init_net, caffe2_predict_net)
    inputs_map = {input[0]:input[1] for input in inputs}

    output = p.run(inputs_map)
    c2_workspace.ResetWorkspace()
    return output 

def run_caffe2_model(predict_net_path, init_net_path, feed_dict):
    from caffe2.python import workspace
    with open(init_net_path, "rb") as f:
        init_net = f.read()
    with open(predict_net_path, "rb") as f:
        predict_net = f.read()

    predictor = workspace.Predictor(init_net, predict_net)
    return [np.array(arr) for arr in predictor.run(feed_dict)] 

def __init__(self, args, config):
        self.args = args
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.cache_dir = utils.get_cache_dir(config)
        self.model_dir = utils.get_model_dir(config)
        _, self.num_parts = utils.get_dataset_mappers(config)
        self.limbs_index = utils.get_limbs_index(config)
        if args.debug is None:
            self.draw_cluster = utils.visualize.DrawCluster(colors=args.colors, thickness=args.thickness)
        else:
            self.draw_feature = utils.visualize.DrawFeature()
            s = re.search('(-?[0-9]+)([a-z]+)(-?[0-9]+)', args.debug)
            stage = int(s.group(1))
            name = s.group(2)
            channel = int(s.group(3))
            self.get_feature = lambda outputs: outputs[stage][name][0][channel]
        self.height, self.width = tuple(map(int, config.get('image', 'size').split()))
        if args.caffe:
            init_net = caffe2_pb2.NetDef()
            with open(os.path.join(self.model_dir, 'init_net.pb'), 'rb') as f:
                init_net.ParseFromString(f.read())
            predict_net = caffe2_pb2.NetDef()
            with open(os.path.join(self.model_dir, 'predict_net.pb'), 'rb') as f:
                predict_net.ParseFromString(f.read())
            p = workspace.Predictor(init_net, predict_net)
            self.inference = lambda tensor: [{'parts': torch.from_numpy(parts), 'limbs': torch.from_numpy(limbs)} for parts, limbs in zip(*[iter(p.run([tensor.detach().cpu().numpy()]))] * 2)]
        else:
            self.step, self.epoch, self.dnn, self.stages = self.load()
            self.inference = model.Inference(config, self.dnn, self.stages)
            self.inference.eval()
            if torch.cuda.is_available():
                self.inference.cuda()
            logging.info(humanize.naturalsize(sum(var.cpu().numpy().nbytes for var in self.inference.state_dict().values())))
        self.cap = self.create_cap()
        self.keys = set(args.keys)
        self.resize = transform.parse_transform(config, config.get('transform', 'resize_test'))
        self.transform_image = transform.get_transform(config, config.get('transform', 'image_test').split())
        self.transform_tensor = transform.get_transform(config, config.get('transform', 'tensor').split())