style_transfer.py 文件源码

def init_model(resp_q, caffe_path, model, weights, mean):
    """Puts the list of layer shapes into resp_q. To be run in a separate process."""
    global logger
    setup_exceptions()
    logger = log_utils.setup_logger('init_model')

    if caffe_path:
        sys.path.append(caffe_path + '/python')
    import caffe
    caffe.set_mode_cpu()
    model = CaffeModel(model, weights, mean)
    shapes = OrderedDict()
    for layer in model.layers():
        shapes[layer] = model.data[layer].shape
    resp_q.put(shapes)