tools.py 文件源码

def _loadModel(self, model_dirs, id):
        print 'loading model...from{}'.format(model_dirs)
        model_file = osp.join(model_dirs, 'vgg16.prototxt')
        model_weights = osp.join(model_dirs, 'vgg16.caffemodel')
        mean_file = osp.join(model_dirs, 'vgg16_mean.npy')
        if id == -1:
            caffe.set_mode_cpu()
        else:
            caffe.set_mode_gpu()
            caffe.set_device(id)
        net = caffe.Net(model_file, model_weights, caffe.TEST)
        transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
        transformer.set_mean('data', np.load(mean_file).mean(1).mean(1))
        transformer.set_channel_swap('data', (2, 1, 0))
        transformer.set_transpose('data', (2, 0, 1))
        #transformer.set_raw_scale('data', 255)
        self.net = net
        self.transformer = transformer
        self.style_layers = VGG16_STYLES
        self.content_layers = VGG16_CONTENTS
        self.layers = VGG16_LAYERS
        print 'model loading done'