mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-09-09 11:15:19 +08:00
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
#!/usr/bin/env python
|
|
|
|
from __future__ import print_function
|
|
|
|
import sys
|
|
from os.path import splitext
|
|
import numpy as np
|
|
|
|
import caffe
|
|
|
|
# Not needed because Tensorflow and Caffe do convolution the same way
|
|
# Needed for conversion to Theano
|
|
|
|
|
|
def rot90(W):
|
|
for i in range(W.shape[0]):
|
|
for j in range(W.shape[1]):
|
|
W[i, j] = np.rot90(W[i, j], 2)
|
|
return W
|
|
|
|
|
|
weights = {}
|
|
assert "prototxt" in splitext(sys.argv[1])[1], "First argument must be caffe prototxt %s" % sys.argv[1]
|
|
assert "caffemodel" in splitext(sys.argv[2])[1], "Second argument must be caffe weights %s" % sys.argv[2]
|
|
net = caffe.Net(sys.argv[1], sys.argv[2], caffe.TEST)
|
|
for k, v in net.params.items():
|
|
print ("Layer %s, has %d params." % (k, len(v)))
|
|
if len(v) == 1:
|
|
W = v[0].data[...]
|
|
W = np.transpose(W, (2, 3, 1, 0))
|
|
weights[k] = {"weights": W}
|
|
elif len(v) == 2:
|
|
W = v[0].data[...]
|
|
W = np.transpose(W, (2, 3, 1, 0))
|
|
b = v[1].data[...]
|
|
weights[k] = {"weights": W, "biases": b}
|
|
elif len(v) == 4: # Batchnorm layer
|
|
k = k.replace('/', '_')
|
|
mean = v[0].data[...]
|
|
variance = v[1].data[...]
|
|
scale = v[2].data[...]
|
|
offset = v[3].data[...]
|
|
weights[k] = {"mean": mean, "variance": variance, "scale": scale, "offset": offset}
|
|
else:
|
|
print("Undefined layer")
|
|
exit()
|
|
|
|
arr = np.asarray(weights)
|
|
weights_name = splitext(sys.argv[2])[0]+".npy"
|
|
np.save(weights_name.lower(), arr)
|