convert.py 文件源码

def load_caffe(img_p, layers=50):
    caffe.set_mode_cpu()

    prototxt = "data/ResNet-%d-deploy.prototxt" % layers
    caffemodel = "data/ResNet-%d-model.caffemodel" % layers
    # net = caffe.Net(prototxt, caffe.TEST)
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)

    net.blobs['data'].data[0] = img_p.transpose((2, 0, 1))
    assert net.blobs['data'].data[0].shape == (3, 224, 224)
    net.forward()

    caffe_prob = net.blobs['prob'].data[0]
    print_prob(caffe_prob)

    return net


# returns the top1 string