feature_extractor.py 文件源码

def __init__(self, path_to_deploy_file, path_to_model_file, input_layer_name="data_q", gpu_mode=True, device_id=1,
                 height=None, width=None):
        self.path_to_deploy_file = path_to_deploy_file
        self.path_to_model_file = path_to_model_file
        if gpu_mode:
            caffe.set_mode_gpu()
            caffe.set_device(device_id)
        else:
            caffe.set_mode_cpu()
        self.net = caffe.Net(path_to_deploy_file, path_to_model_file, caffe.TEST)
        self.input_layer_name = input_layer_name
        self.height = height or self.net.blobs[self.input_layer_name].data.shape[2]
        self.width = width or self.net.blobs[self.input_layer_name].data.shape[3]