diff --git a/layers_builder.py b/layers_builder.py index 724f79c..e007117 100644 --- a/layers_builder.py +++ b/layers_builder.py @@ -1,3 +1,5 @@ +from __future__ import print_function +from math import ceil from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D from keras.layers import BatchNormalization, Activation, Input, Dropout, ZeroPadding2D, Lambda from keras.layers.merge import Concatenate, Add @@ -6,137 +8,153 @@ from keras.optimizers import SGD import tensorflow as tf -learning_rate = 1e-3 # Layer specific learning rate +learning_rate = 1e-3 # Layer specific learning rate # Weight decay not implemented + def BN(name=""): return BatchNormalization(momentum=0.95, name=name, epsilon=1e-5) -def Interp(x, shape=(60,60)): - new_height,new_width = shape - resized = tf.image.resize_images(x, [new_height, new_width], align_corners=True) + +def Interp(x, shape): + from keras.backend import tf as ktf + new_height, new_width = shape + resized = ktf.image.resize_images(x, [new_height, new_width], align_corners=True) return resized + def residual_conv(prev, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False): lvl = str(lvl) sub_lvl = str(sub_lvl) - names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_reduce" , - "conv"+lvl+"_"+ sub_lvl +"_1x1_reduce_bn", - "conv"+lvl+"_"+ sub_lvl +"_3x3", - "conv"+lvl+"_"+ sub_lvl +"_3x3_bn", - "conv"+lvl+"_"+ sub_lvl +"_1x1_increase", - "conv"+lvl+"_"+ sub_lvl +"_1x1_increase_bn"] - if modify_stride == False: - prev = Conv2D(64 * level, (1,1), strides=(1,1), name=names[0], use_bias=False)(prev) - elif modify_stride == True: - prev = Conv2D(64 * level, (1,1), strides=(2,2), name=names[0], use_bias=False)(prev) + names = ["conv"+lvl+"_" + sub_lvl + "_1x1_reduce", + "conv"+lvl+"_" + sub_lvl + "_1x1_reduce_bn", + "conv"+lvl+"_" + sub_lvl + "_3x3", + "conv"+lvl+"_" + sub_lvl + "_3x3_bn", + "conv"+lvl+"_" + sub_lvl + "_1x1_increase", + "conv"+lvl+"_" + sub_lvl + "_1x1_increase_bn"] + if modify_stride is False: + prev = Conv2D(64 * level, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev) + elif modify_stride is True: + prev = Conv2D(64 * level, (1, 1), strides=(2, 2), name=names[0], use_bias=False)(prev) prev = BN(name=names[1])(prev) prev = Activation('relu')(prev) - prev = ZeroPadding2D(padding=(pad,pad))(prev) - prev = Conv2D(64 * level, (3,3), strides=(1,1), dilation_rate=pad, name=names[2], use_bias=False)(prev) + prev = ZeroPadding2D(padding=(pad, pad))(prev) + prev = Conv2D(64 * level, (3, 3), strides=(1, 1), dilation_rate=pad, name=names[2], use_bias=False)(prev) prev = BN(name=names[3])(prev) prev = Activation('relu')(prev) - prev = Conv2D(256 * level, (1,1), strides=(1,1), name=names[4], use_bias=False)(prev) + prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[4], use_bias=False)(prev) prev = BN(name=names[5])(prev) return prev + def short_convolution_branch(prev, level, lvl=1, sub_lvl=1, modify_stride=False): lvl = str(lvl) sub_lvl = str(sub_lvl) - names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_proj", - "conv"+lvl+"_"+ sub_lvl +"_1x1_proj_bn"] + names = ["conv" + lvl+"_" + sub_lvl + "_1x1_proj", + "conv" + lvl+"_" + sub_lvl + "_1x1_proj_bn"] - if modify_stride == False: - prev = Conv2D(256 * level ,(1,1), strides=(1,1), name=names[0], use_bias=False)(prev) - elif modify_stride == True: - prev = Conv2D(256 * level, (1,1), strides=(2,2), name=names[0], use_bias=False)(prev) + if modify_stride is False: + prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev) + elif modify_stride is True: + prev = Conv2D(256 * level, (1, 1), strides=(2, 2), name=names[0], use_bias=False)(prev) prev = BN(name=names[1])(prev) return prev + def empty_branch(prev): return prev + def residual_short(prev_layer, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False): prev_layer = Activation('relu')(prev_layer) block_1 = residual_conv(prev_layer, level, - pad=pad, lvl=lvl, sub_lvl=sub_lvl, - modify_stride=modify_stride) + pad=pad, lvl=lvl, sub_lvl=sub_lvl, + modify_stride=modify_stride) block_2 = short_convolution_branch(prev_layer, level, - lvl=lvl, sub_lvl=sub_lvl, - modify_stride=modify_stride) + lvl=lvl, sub_lvl=sub_lvl, + modify_stride=modify_stride) added = Add()([block_1, block_2]) return added + def residual_empty(prev_layer, level, pad=1, lvl=1, sub_lvl=1): prev_layer = Activation('relu')(prev_layer) - block_1 = residual_conv(prev_layer, level, - pad=pad, lvl=lvl, sub_lvl=sub_lvl) + block_1 = residual_conv(prev_layer, level, pad=pad, lvl=lvl, sub_lvl=sub_lvl) block_2 = empty_branch(prev_layer) added = Add()([block_1, block_2]) return added -def ResNet(inp): - #Names for the first couple layers of model + +def ResNet(inp, layers): + # Names for the first couple layers of model names = ["conv1_1_3x3_s2", - "conv1_1_3x3_s2_bn", - "conv1_2_3x3", - "conv1_2_3x3_bn", - "conv1_3_3x3", - "conv1_3_3x3_bn"] + "conv1_1_3x3_s2_bn", + "conv1_2_3x3", + "conv1_2_3x3_bn", + "conv1_3_3x3", + "conv1_3_3x3_bn"] - #---Short branch(only start of network) + # Short branch(only start of network) - cnv1 = Conv2D(64, (3, 3), strides=(2, 2), padding='same', name=names[0], use_bias=False)(inp) # "conv1_1_3x3_s2" + cnv1 = Conv2D(64, (3, 3), strides=(2, 2), padding='same', name=names[0], use_bias=False)(inp) # "conv1_1_3x3_s2" bn1 = BN(name=names[1])(cnv1) # "conv1_1_3x3_s2/bn" - relu1 = Activation('relu')(bn1) #"conv1_1_3x3_s2/relu" + relu1 = Activation('relu')(bn1) # "conv1_1_3x3_s2/relu" - cnv1 = Conv2D(64, (3, 3), strides=(1, 1), padding='same', name=names[2], use_bias=False)(relu1) #"conv1_2_3x3" - bn1 = BN(name=names[3])(cnv1) #"conv1_2_3x3/bn" - relu1 = Activation('relu')(bn1) #"conv1_2_3x3/relu" + cnv1 = Conv2D(64, (3, 3), strides=(1, 1), padding='same', name=names[2], use_bias=False)(relu1) # "conv1_2_3x3" + bn1 = BN(name=names[3])(cnv1) # "conv1_2_3x3/bn" + relu1 = Activation('relu')(bn1) # "conv1_2_3x3/relu" - cnv1 = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name=names[4], use_bias=False)(relu1) #"conv1_3_3x3" - bn1 = BN(name=names[5])(cnv1) #"conv1_3_3x3/bn" - relu1 = Activation('relu')(bn1) #"conv1_3_3x3/relu" + cnv1 = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name=names[4], use_bias=False)(relu1) # "conv1_3_3x3" + bn1 = BN(name=names[5])(cnv1) # "conv1_3_3x3/bn" + relu1 = Activation('relu')(bn1) # "conv1_3_3x3/relu" - res = MaxPooling2D(pool_size=(3,3), padding='same', strides=(2,2))(relu1) #"pool1_3x3_s2" - - #---Residual layers(body of network) + res = MaxPooling2D(pool_size=(3, 3), padding='same', strides=(2, 2))(relu1) # "pool1_3x3_s2" + + # ---Residual layers(body of network) """ Modify_stride --Used only once in first 3_1 convolutions block. changes stride of first convolution from 1 -> 2 """ - #2_1- 2_3 - res = residual_short(res, 1, pad=1, lvl=2, sub_lvl=1) + # 2_1- 2_3 + res = residual_short(res, 1, pad=1, lvl=2, sub_lvl=1) for i in range(2): - res = residual_empty(res, 1, pad=1, lvl=2, sub_lvl=i+2) + res = residual_empty(res, 1, pad=1, lvl=2, sub_lvl=i+2) - #3_1 - 3_3 - res = residual_short(res, 2, pad=1, lvl=3, sub_lvl=1, modify_stride=True) + # 3_1 - 3_3 + res = residual_short(res, 2, pad=1, lvl=3, sub_lvl=1, modify_stride=True) for i in range(3): - res = residual_empty(res, 2, pad=1, lvl=3, sub_lvl=i+2) + res = residual_empty(res, 2, pad=1, lvl=3, sub_lvl=i+2) + if layers is 50: + # 4_1 - 4_6 + res = residual_short(res, 4, pad=2, lvl=4, sub_lvl=1) + for i in range(5): + res = residual_empty(res, 4, pad=2, lvl=4, sub_lvl=i+2) + elif layers is 101: + # 4_1 - 4_23 + res = residual_short(res, 4, pad=2, lvl=4, sub_lvl=1) + for i in range(22): + res = residual_empty(res, 4, pad=2, lvl=4, sub_lvl=i+2) + else: + print("This ResNet is not implemented") - #4_1 - 4_6 - res = residual_short(res, 4, pad=2, lvl=4, sub_lvl=1) - for i in range(5): - res = residual_empty(res, 4, pad=2, lvl=4, sub_lvl=i+2) - - #5_1 - 5_3 - res = residual_short(res, 8, pad=4, lvl=5, sub_lvl=1) + # 5_1 - 5_3 + res = residual_short(res, 8, pad=4, lvl=5, sub_lvl=1) for i in range(2): res = residual_empty(res, 8, pad=4, lvl=5, sub_lvl=i+2) res = Activation('relu')(res) return res -def interp_block(prev_layer, level, str_lvl=1): + +def interp_block(prev_layer, level, feature_map_shape, str_lvl=1, ): str_lvl = str(str_lvl) @@ -147,45 +165,49 @@ def interp_block(prev_layer, level, str_lvl=1): kernel = (10*level, 10*level) strides = (10*level, 10*level) - prev_layer = AveragePooling2D(kernel,strides=strides)(prev_layer) - prev_layer = Conv2D(512, (1,1), strides=(1,1), name=names[0], use_bias=False)(prev_layer) + prev_layer = AveragePooling2D(kernel, strides=strides)(prev_layer) + prev_layer = Conv2D(512, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev_layer) prev_layer = BN(name=names[1])(prev_layer) prev_layer = Activation('relu')(prev_layer) - prev_layer = Lambda(Interp)(prev_layer) + prev_layer = Lambda(Interp, arguments={'shape': feature_map_shape})(prev_layer) return prev_layer -def PSPNet(res): - #---PSPNet concat layers with Interpolation +def PSPNet(res, input_shape): + """Build the Pyramid Pooling Module.""" + # ---PSPNet concat layers with Interpolation + feature_map_size = tuple(int(ceil(input_dim / 8.0)) for input_dim in input_shape) + print("PSP module will interpolate to a final feature map size of %s" % (feature_map_size, )) - interp_block1 = interp_block(res, 6, str_lvl=1) - interp_block2 = interp_block(res, 3, str_lvl=2) - interp_block3 = interp_block(res, 2, str_lvl=3) - interp_block6 = interp_block(res, 1, str_lvl=6) + interp_block1 = interp_block(res, 6, feature_map_size, str_lvl=1) + interp_block2 = interp_block(res, 3, feature_map_size, str_lvl=2) + interp_block3 = interp_block(res, 2, feature_map_size, str_lvl=3) + interp_block6 = interp_block(res, 1, feature_map_size, str_lvl=6) - #concat all these layers. resulted shape=(1,60,60,4096) + # concat all these layers. resulted shape=(1,60,60,4096) res = Concatenate()([res, - interp_block6, - interp_block3, - interp_block2, - interp_block1]) + interp_block6, + interp_block3, + interp_block2, + interp_block1]) return res -def build_pspnet(activation='softmax'): - ''' - Normal PSPNet. - ''' - inp = Input((473,473,3)) - res = ResNet(inp) - psp = PSPNet(res) + +def build_pspnet(nb_classes, resnet_layers, input_shape, activation='softmax'): + """Build PSPNet.""" + print("Building a PSPNet based on ResNet %i expecting inputs of shape %s predicting %i classes" % (resnet_layers, input_shape, nb_classes)) + + inp = Input((input_shape[0], input_shape[1], 3)) + res = ResNet(inp, layers=resnet_layers) + psp = PSPNet(res, input_shape) x = Conv2D(512, (3, 3), strides=(1, 1), padding="same", name="conv5_4", use_bias=False)(psp) x = BN(name="conv5_4_bn")(x) x = Activation('relu')(x) x = Dropout(0.1)(x) - x = Conv2D(150, (1, 1), strides=(1, 1), name="conv6")(x) - x = Lambda(Interp, arguments={'shape': (473,473)})(x) + x = Conv2D(nb_classes, (1, 1), strides=(1, 1), name="conv6")(x) + x = Lambda(Interp, arguments={'shape': (input_shape[0], input_shape[1])})(x) x = Activation('softmax')(x) model = Model(inputs=inp, outputs=x) @@ -193,6 +215,6 @@ def build_pspnet(activation='softmax'): # Solver sgd = SGD(lr=learning_rate, momentum=0.9, nesterov=True) model.compile(optimizer=sgd, - loss='categorical_crossentropy', - metrics=['accuracy']) + loss='categorical_crossentropy', + metrics=['accuracy']) return model diff --git a/pspnet.py b/pspnet.py index fa42f64..1bfa6ed 100644 --- a/pspnet.py +++ b/pspnet.py @@ -1,79 +1,122 @@ +from __future__ import print_function import os +from os.path import splitext import argparse import numpy as np from scipy import misc, ndimage - from keras import backend as K +from keras.models import model_from_json import tensorflow as tf - import layers_builder as layers import utils -WEIGHTS = 'pspnet50_ade20k.npy' -DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB -class PSPNet: +DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB, these are the means for the ImageNet pretrained ResNet - def __init__(self): - self.model = layers.build_pspnet() - set_npy_weights(self.model, WEIGHTS) + +class PSPNet(object): + """Pyramid Scene Parsing Network by Hengshuang Zhao et al 2017""" + + def __init__(self, nb_classes, resnet_layers, input_shape, weights): + self.input_shape = input_shape + json_path = weights + ".json" + h5_path = weights + ".h5" + if os.path.isfile(json_path) and os.path.isfile(h5_path): + print("Keras model & weights found, loading...") + with open(json_path, 'r') as file_handle: + self.model = model_from_json(file_handle.read()) + self.model.load_weights(h5_path) + else: + print("No Keras model & weights found, importing from numpy weights.") + self.model = layers.build_pspnet(nb_classes=nb_classes, resnet_layers=resnet_layers, input_shape=self.input_shape) + self.set_npy_weights(weights) def predict(self, img): - ''' + """ + Predict segementation for an image. + Arguments: - img: must be 473x473x3 - ''' - h_ori,w_ori = img.shape[:2] + img: must be rowsxcolsx3 + """ + h_ori, w_ori = img.shape[:2] # Preprocess - img = misc.imresize(img, (473, 473)) + img = misc.imresize(img, self.input_shape) + img = img - DATA_MEAN - img = img[:,:,::-1] # RGB => BGR + img = img[:, :, ::-1] # RGB => BGR img = img.astype('float32') + print("Predicting...") probs = self.feed_forward(img) - h,w = probs.shape[:2] - probs = ndimage.zoom(probs, (1.*h_ori/h,1.*w_ori/w,1.), order=1, prefilter=False) + h, w = probs.shape[:2] + probs = ndimage.zoom(probs, (1.*h_ori/h, 1.*w_ori/w, 1.), order=1, prefilter=False) + print("Finished prediction...") + return probs - def predict_sliding_window(self, img): - pass - def feed_forward(self, data): - assert data.shape == (473,473,3) - data = data[np.newaxis,:,:,:] + assert data.shape == (self.input_shape[0], self.input_shape[1], 3) + data = data[np.newaxis, :, :, :] # utils.debug(self.model, data) pred = self.model.predict(data) return pred[0] -def set_npy_weights(model, npy_weights): - weights = np.load(npy_weights).item() + def set_npy_weights(self, weights_path): + npy_weights_path = weights_path + ".npy" + json_path = weights_path + ".json" + h5_path = weights_path + ".h5" - for layer in model.layers: - print layer.name - if layer.name[:4] == 'conv' and layer.name[-2:] == 'bn': - mean = weights[layer.name]['mean'].reshape(-1) - variance = weights[layer.name]['variance'].reshape(-1) - scale = weights[layer.name]['scale'].reshape(-1) - offset = weights[layer.name]['offset'].reshape(-1) - - model.get_layer(layer.name).set_weights([mean, variance, scale, offset]) + print("Importing weights from %s" % npy_weights_path) + weights = np.load(npy_weights_path).item() + + for layer in self.model.layers: + print(layer.name) + if layer.name[:4] == 'conv' and layer.name[-2:] == 'bn': + mean = weights[layer.name]['mean'].reshape(-1) + variance = weights[layer.name]['variance'].reshape(-1) + scale = weights[layer.name]['scale'].reshape(-1) + offset = weights[layer.name]['offset'].reshape(-1) + + self.model.get_layer(layer.name).set_weights([mean, variance, scale, offset]) + + elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu': + try: + weight = weights[layer.name]['weights'] + self.model.get_layer(layer.name).set_weights([weight]) + except Exception as err: + biases = weights[layer.name]['biases'] + self.model.get_layer(layer.name).set_weights([weight, biases]) + print('Finished importing weights.') + + print("Writing keras model & weights") + json_string = self.model.to_json() + with open(json_path, 'w') as file_handle: + file_handle.write(json_string) + self.model.save_weights(h5_path) + print("Finished writing Keras model & weights") + + +class PSPNet50(PSPNet): + """Build a PSPNet based on a 50-Layer ResNet.""" + + def __init__(self, nb_classes, weights, input_shape): + PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=50, input_shape=input_shape, weights=weights) + + +class PSPNet101(PSPNet): + """Build a PSPNet based on a 101-Layer ResNet.""" + + def __init__(self, nb_classes, weights, input_shape): + PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=101, input_shape=input_shape, weights=weights) - elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu': - try: - weight = weights[layer.name]['weights'] - model.get_layer(layer.name).set_weights([weight]) - except Exception as err: - biases = weights[layer.name]['biases'] - model.get_layer(layer.name).set_weights([weight, biases]) - print 'Finished.' - return model if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--input_path', type=str, default='', required=True, help='Path the input image') - parser.add_argument('--output_path', type=str, default='', required=True, help='Path to output') + parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k', help='Model/Weights to use', choices=['pspnet50_ade20k', 'pspnet101_cityscapes', 'pspnet101_voc2012']) + parser.add_argument('-i', '--input_path', type=str, default='test.jpg', help='Path the input image') + parser.add_argument('-o', '--output_path', type=str, default='test.jpg', help='Path to output') parser.add_argument('--id', default="0") args = parser.parse_args() @@ -84,13 +127,27 @@ if __name__ == "__main__": with sess.as_default(): img = misc.imread(args.input_path) - - pspnet = PSPNet() + print(args) + + if "pspnet50" in args.model: + pspnet = PSPNet50(nb_classes=150, input_shape=(473, 473), weights=args.model) + elif "pspnet101" in args.model: + if "cityscapes" in args.model: + pspnet = PSPNet101(nb_classes=19, input_shape=(713, 713), weights=args.model) + if "voc2012" in args.model: + pspnet = PSPNet101(nb_classes=21, input_shape=(473, 473), weights=args.model) + + else: + print("Network architecture not implemented.") + probs = pspnet.predict(img) + print("Writing results...") cm = np.argmax(probs, axis=2) + 1 pm = np.max(probs, axis=2) color_cm = utils.add_color(cm) - misc.imsave(args.output_path, color_cm) - misc.imsave("probs.jpg", pm) - + alpha_blended = 0.5 * color_cm * 255 + 0.5 * img # color cm is [0.0-1.0] img [0-255] + filename, ext = splitext(args.output_path) + misc.imsave(filename + "_seg" + ext, color_cm) + misc.imsave(filename + "_probs" + ext, pm) + misc.imsave(filename + "_seg_blended" + ext, alpha_blended) diff --git a/utils.py b/utils.py index bbf9d82..8f179f5 100644 --- a/utils.py +++ b/utils.py @@ -1,20 +1,23 @@ +from __future__ import print_function import colorsys import numpy as np from keras.models import Model + def add_color(img): - h,w = img.shape - img_color = np.zeros((h,w,3)) - for i in xrange(1,151): + h, w = img.shape + img_color = np.zeros((h, w, 3)) + for i in xrange(1, 151): img_color[img == i] = to_color(i) return img_color + def to_color(category): # Maps each category a good distance away # from each other on the HSV color space v = (category-1)*(137.5/360) - return colorsys.hsv_to_rgb(v,1,1) + return colorsys.hsv_to_rgb(v, 1, 1) # For printing the activations in each layer @@ -23,10 +26,14 @@ def debug(model, data): names = [layer.name for layer in model.layers] for name in names[:]: print_activation(model, name, data) + + def print_activation(model, layer_name, data): intermediate_layer_model = Model(inputs=model.input, outputs=model.get_layer(layer_name).output) io = intermediate_layer_model.predict(data) - print layer_name, array_to_str(io) + print (layer_name, array_to_str(io)) + + def array_to_str(a): - return "{} {} {} {} {}".format(a.dtype, a.shape, np.min(a), np.max(a), np.mean(a)) \ No newline at end of file + return "{} {} {} {} {}".format(a.dtype, a.shape, np.min(a), np.max(a), np.mean(a)) diff --git a/weight_converter.py b/weight_converter.py index 033beb1..e98ffcd 100644 --- a/weight_converter.py +++ b/weight_converter.py @@ -1,42 +1,48 @@ -import os +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: - 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() +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: + 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) -np.save("pspnet50_ade20k.npy", arr) - - +weights_name = splitext(sys.argv[2])[0]+".npy" +np.save(weights_name.lower(), arr)