def load_pretrained_model(self, model_path):
"""
Load weights from caffemodel w/o caffe dependency
and plug them in corresponding modules
"""
# My eyes and my heart both hurt when writing this method
# Only care about layer_types that have trainable parameters
ltypes = [
"BNData",
"ConvolutionData",
"HoleConvolutionData",
"Convolution",
] # Convolution type for conv3_sub1_proj
def _get_layer_params(layer, ltype):
if ltype == "BNData":
gamma = np.array(layer.blobs[0].data)
beta = np.array(layer.blobs[1].data)
mean = np.array(layer.blobs[2].data)
var = np.array(layer.blobs[3].data)
return [mean, var, gamma, beta]
elif ltype in [
"ConvolutionData", "HoleConvolutionData", "Convolution"
]:
is_bias = layer.convolution_param.bias_term
weights = np.array(layer.blobs[0].data)
bias = []
if is_bias:
bias = np.array(layer.blobs[1].data)
return [weights, bias]
elif ltype == "InnerProduct":
raise Exception(
"Fully connected layers {}, not supported".format(ltype))
else:
raise Exception("Unkown layer type {}".format(ltype))
net = caffe_pb2.NetParameter()
with open(model_path, "rb") as model_file:
net.MergeFromString(model_file.read())
# dict formatted as -> key:<layer_name> :: value:<layer_type>
layer_types = {}
# dict formatted as -> key:<layer_name> :: value:[<list_of_params>]
layer_params = {}
for l in net.layer:
lname = l.name
ltype = l.type
lbottom = l.bottom
ltop = l.top
if ltype in ltypes:
print("Processing layer {} | {}, {}".format(
lname, lbottom, ltop))
layer_types[lname] = ltype
layer_params[lname] = _get_layer_params(l, ltype)
# if len(l.blobs) > 0:
# print(lname, ltype, lbottom, ltop, len(l.blobs))
# Set affine=False for all batchnorm modules
def _no_affine_bn(module=None):
if isinstance(module, nn.BatchNorm2d):
module.affine = False
if len([m for m in module.children()]) > 0:
for child in module.children():
_no_affine_bn(child)
# _no_affine_bn(self)
def _transfer_conv(layer_name, module):
weights, bias = layer_params[layer_name]
w_shape = np.array(module.weight.size())
print("CONV {}: Original {} and trans weights {}".format(
layer_name, w_shape, weights.shape))
module.weight.data.copy_(
torch.from_numpy(weights).view_as(module.weight))
if len(bias) != 0:
b_shape = np.array(module.bias.size())
print("CONV {}: Original {} and trans bias {}".format(
layer_name, b_shape, bias.shape))
module.bias.data.copy_(
torch.from_numpy(bias).view_as(module.bias))
def _transfer_bn(conv_layer_name, bn_module):
mean, var, gamma, beta = layer_params[conv_layer_name + "/bn"]
print("BN {}: Original {} and trans weights {}".format(
conv_layer_name, bn_module.running_mean.size(), mean.shape))
bn_module.running_mean.copy_(
torch.from_numpy(mean).view_as(bn_module.running_mean))
bn_module.running_var.copy_(
torch.from_numpy(var).view_as(bn_module.running_var))
bn_module.weight.data.copy_(
torch.from_numpy(gamma).view_as(bn_module.weight))
bn_module.bias.data.copy_(
torch.from_numpy(beta).view_as(bn_module.bias))
def _transfer_conv_bn(conv_layer_name, mother_module):
conv_module = mother_module[0]
_transfer_conv(conv_layer_name, conv_module)
if conv_layer_name + "/bn" in layer_params.keys():
bn_module = mother_module[1]
_transfer_bn(conv_layer_name, bn_module)
def _transfer_residual(block_name, block):
block_module, n_layers = block[0], block[1]
prefix = block_name[:5]
if ("bottleneck"
in block_name) or ("identity"
not in block_name): # Conv block
bottleneck = block_module.layers[0]
bottleneck_conv_bn_dic = {
prefix + "_1_1x1_reduce": bottleneck.cbr1.cbr_unit,
prefix + "_1_3x3": bottleneck.cbr2.cbr_unit,
prefix + "_1_1x1_proj": bottleneck.cb4.cb_unit,
prefix + "_1_1x1_increase": bottleneck.cb3.cb_unit,
}
for k, v in bottleneck_conv_bn_dic.items():
_transfer_conv_bn(k, v)
if ("identity"
in block_name) or ("bottleneck"
not in block_name): # Identity blocks
base_idx = 2 if "identity" in block_name else 1
for layer_idx in range(2, n_layers + 1):
residual_layer = block_module.layers[layer_idx - base_idx]
residual_conv_bn_dic = {
"_".join(map(str, [prefix, layer_idx, "1x1_reduce"])):
residual_layer.cbr1.cbr_unit,
"_".join(map(str, [prefix, layer_idx, "3x3"])):
residual_layer.cbr2.cbr_unit,
"_".join(map(str, [prefix, layer_idx, "1x1_increase"])):
residual_layer.cb3.cb_unit,
}
for k, v in residual_conv_bn_dic.items():
_transfer_conv_bn(k, v)
convbn_layer_mapping = {
"conv1_1_3x3_s2": self.convbnrelu1_1.cbr_unit,
"conv1_2_3x3": self.convbnrelu1_2.cbr_unit,
"conv1_3_3x3": self.convbnrelu1_3.cbr_unit,
"conv1_sub1": self.convbnrelu1_sub1.cbr_unit,
"conv2_sub1": self.convbnrelu2_sub1.cbr_unit,
"conv3_sub1": self.convbnrelu3_sub1.cbr_unit,
# 'conv5_3_pool6_conv': self.pyramid_pooling.paths[0].cbr_unit,
# 'conv5_3_pool3_conv': self.pyramid_pooling.paths[1].cbr_unit,
# 'conv5_3_pool2_conv': self.pyramid_pooling.paths[2].cbr_unit,
# 'conv5_3_pool1_conv': self.pyramid_pooling.paths[3].cbr_unit,
"conv5_4_k1": self.conv5_4_k1.cbr_unit,
"conv_sub4": self.cff_sub24.low_dilated_conv_bn.cb_unit,
"conv3_1_sub2_proj": self.cff_sub24.high_proj_conv_bn.cb_unit,
"conv_sub2": self.cff_sub12.low_dilated_conv_bn.cb_unit,
"conv3_sub1_proj": self.cff_sub12.high_proj_conv_bn.cb_unit,
}
residual_layers = {
"conv2": [self.res_block2, self.block_config[0]],
"conv3_bottleneck": [self.res_block3_conv, self.block_config[1]],
"conv3_identity": [self.res_block3_identity, self.block_config[1]],
"conv4": [self.res_block4, self.block_config[2]],
"conv5": [self.res_block5, self.block_config[3]],
}
# Transfer weights for all non-residual conv+bn layers
for k, v in convbn_layer_mapping.items():
_transfer_conv_bn(k, v)
# Transfer weights for final non-bn conv layer
_transfer_conv("conv6_cls", self.classification)
_transfer_conv("conv6_sub4", self.cff_sub24.low_classifier_conv)
_transfer_conv("conv6_sub2", self.cff_sub12.low_classifier_conv)
# Transfer weights for all residual layers
for k, v in residual_layers.items():
_transfer_residual(k, v)
def load_pretrained_model(self, model_path):
"""
Load weights from caffemodel w/o caffe dependency
and plug them in corresponding modules
"""
# My eyes and my heart both hurt when writing this method
# Only care about layer_types that have trainable parameters
ltypes = ['BNData', 'ConvolutionData', 'HoleConvolutionData']
def _get_layer_params(layer, ltype):
if ltype == 'BNData':
gamma = np.array(layer.blobs[0].data)
beta = np.array(layer.blobs[1].data)
mean = np.array(layer.blobs[2].data)
var = np.array(layer.blobs[3].data)
return [mean, var, gamma, beta]
elif ltype in ['ConvolutionData', 'HoleConvolutionData']:
is_bias = layer.convolution_param.bias_term
weights = np.array(layer.blobs[0].data)
bias = []
if is_bias:
bias = np.array(layer.blobs[1].data)
return [weights, bias]
elif ltype == 'InnerProduct':
raise Exception(
"Fully connected layers {}, not supported".format(ltype))
else:
raise Exception("Unkown layer type {}".format(ltype))
net = caffe_pb2.NetParameter()
with open(model_path, 'rb') as model_file:
net.MergeFromString(model_file.read())
# dict formatted as -> key:<layer_name> :: value:<layer_type>
layer_types = {}
# dict formatted as -> key:<layer_name> :: value:[<list_of_params>]
layer_params = {}
for l in net.layer:
lname = l.name
ltype = l.type
if ltype in ltypes:
print("Processing layer {}".format(lname))
layer_types[lname] = ltype
layer_params[lname] = _get_layer_params(l, ltype)
# Set affine=False for all batchnorm modules
def _no_affine_bn(module=None):
if isinstance(module, nn.BatchNorm2d):
module.affine = False
if len([m for m in module.children()]) > 0:
for child in module.children():
_no_affine_bn(child)
#_no_affine_bn(self)
def _transfer_conv(layer_name, module):
weights, bias = layer_params[layer_name]
w_shape = np.array(module.weight.size())
print("CONV {}: Original {} and trans weights {}".format(
layer_name, w_shape, weights.shape))
module.weight.data.copy_(
torch.from_numpy(weights).view_as(module.weight))
if len(bias) != 0:
b_shape = np.array(module.bias.size())
print("CONV {}: Original {} and trans bias {}".format(
layer_name, b_shape, bias.shape))
module.bias.data.copy_(
torch.from_numpy(bias).view_as(module.bias))
def _transfer_conv_bn(conv_layer_name, mother_module):
conv_module = mother_module[0]
bn_module = mother_module[1]
_transfer_conv(conv_layer_name, conv_module)
mean, var, gamma, beta = layer_params[conv_layer_name + '/bn']
print("BN {}: Original {} and trans weights {}".format(
conv_layer_name, bn_module.running_mean.size(), mean.shape))
bn_module.running_mean.copy_(
torch.from_numpy(mean).view_as(bn_module.running_mean))
bn_module.running_var.copy_(
torch.from_numpy(var).view_as(bn_module.running_var))
bn_module.weight.data.copy_(
torch.from_numpy(gamma).view_as(bn_module.weight))
bn_module.bias.data.copy_(
torch.from_numpy(beta).view_as(bn_module.bias))
def _transfer_residual(prefix, block):
block_module, n_layers = block[0], block[1]
bottleneck = block_module.layers[0]
bottleneck_conv_bn_dic = {
prefix + '_1_1x1_reduce': bottleneck.cbr1.cbr_unit,
prefix + '_1_3x3': bottleneck.cbr2.cbr_unit,
prefix + '_1_1x1_proj': bottleneck.cb4.cb_unit,
prefix + '_1_1x1_increase': bottleneck.cb3.cb_unit,
}
for k, v in bottleneck_conv_bn_dic.items():
_transfer_conv_bn(k, v)
for layer_idx in range(2, n_layers + 1):
residual_layer = block_module.layers[layer_idx - 1]
residual_conv_bn_dic = {
'_'.join(map(str, [prefix, layer_idx, '1x1_reduce'])):
residual_layer.cbr1.cbr_unit,
'_'.join(map(str, [prefix, layer_idx, '3x3'])):
residual_layer.cbr2.cbr_unit,
'_'.join(map(str, [prefix, layer_idx, '1x1_increase'])):
residual_layer.cb3.cb_unit,
}
for k, v in residual_conv_bn_dic.items():
_transfer_conv_bn(k, v)
convbn_layer_mapping = {
'conv1_1_3x3_s2':
self.convbnrelu1_1.cbr_unit,
'conv1_2_3x3':
self.convbnrelu1_2.cbr_unit,
'conv1_3_3x3':
self.convbnrelu1_3.cbr_unit,
'conv5_3_pool6_conv':
self.pyramid_pooling.paths[0].cbr_unit,
'conv5_3_pool3_conv':
self.pyramid_pooling.paths[1].cbr_unit,
'conv5_3_pool2_conv':
self.pyramid_pooling.paths[2].cbr_unit,
'conv5_3_pool1_conv':
self.pyramid_pooling.paths[3].cbr_unit,
'conv5_4':
self.cbr_final.cbr_unit,
'conv4_' + str(self.block_config[2] + 1):
self.convbnrelu4_aux.cbr_unit,
} # Auxiliary layers for training
residual_layers = {
'conv2': [self.res_block2, self.block_config[0]],
'conv3': [self.res_block3, self.block_config[1]],
'conv4': [self.res_block4, self.block_config[2]],
'conv5': [self.res_block5, self.block_config[3]],
}
# Transfer weights for all non-residual conv+bn layers
for k, v in convbn_layer_mapping.items():
_transfer_conv_bn(k, v)
# Transfer weights for final non-bn conv layer
_transfer_conv('conv6', self.classification)
_transfer_conv('conv6_1', self.aux_cls)
# Transfer weights for all residual layers
for k, v in residual_layers.items():
_transfer_residual(k, v)
def load_pretrained_model(self, model_path):
"""
Load weights from caffemodel w/o caffe dependency
and plug them in corresponding modules
"""
# My eyes and my heart both hurt when writing this method
# Only care about layer_types that have trainable parameters
ltypes = ['BNData', 'ConvolutionData', 'HoleConvolutionData']
def _get_layer_params(layer, ltype):
if ltype == 'BNData':
n_channels = layer.blobs[0].shape.dim[1]
gamma = np.array([w for w in layer.blobs[0].data]).reshape(n_channels)
beta = np.array([w for w in layer.blobs[1].data]).reshape(n_channels)
mean = np.array([w for w in layer.blobs[2].data]).reshape(n_channels)
var = np.array([w for w in layer.blobs[3].data]).reshape(n_channels)
return [mean, var, gamma, beta]
elif ltype in ['ConvolutionData', 'HoleConvolutionData']:
is_bias = layer.convolution_param.bias_term
shape = [int(d) for d in layer.blobs[0].shape.dim]
weights = np.array([w for w in layer.blobs[0].data]).reshape(shape)
bias = []
if is_bias:
bias = np.array([w for w in layer.blobs[1].data]).reshape(shape[0])
return [weights, bias]
elif ltype == 'InnerProduct':
raise Exception("Fully connected layers {}, not supported".format(ltype))
else:
raise Exception("Unkown layer type {}".format(ltype))
net = caffe_pb2.NetParameter()
with open(model_path, 'rb') as model_file:
net.MergeFromString(model_file.read())
# dict formatted as -> key:<layer_name> :: value:<layer_type>
layer_types = {}
# dict formatted as -> key:<layer_name> :: value:[<list_of_params>]
layer_params = {}
for l in net.layer:
lname = l.name
ltype = l.type
if ltype in ltypes:
print("Processing layer {}".format(lname))
layer_types[lname] = ltype
layer_params[lname] = _get_layer_params(l, ltype)
# Set affine=False for all batchnorm modules
def _no_affine_bn(module=None):
if isinstance(module, nn.BatchNorm2d):
module.affine = False
if len([m for m in module.children()]) > 0:
for child in module.children():
_no_affine_bn(child)
#_no_affine_bn(self)
def _transfer_conv(layer_name, module):
weights, bias = layer_params[layer_name]
w_shape = np.array(module.weight.size())
np.testing.assert_array_equal(weights.shape, w_shape)
print("CONV: Original {} and trans weights {}".format(w_shape,
weights.shape))
module.weight.data = torch.from_numpy(weights)
if len(bias) != 0:
b_shape = np.array(module.bias.size())
np.testing.assert_array_equal(bias.shape, b_shape)
print("CONV: Original {} and trans bias {}".format(b_shape,
bias.shape))
module.bias.data = torch.from_numpy(bias)
def _transfer_conv_bn(conv_layer_name, mother_module):
conv_module = mother_module[0]
bn_module = mother_module[1]
_transfer_conv(conv_layer_name, conv_module)
mean, var, gamma, beta = layer_params[conv_layer_name+'/bn']
print("BN: Original {} and trans weights {}".format(bn_module.running_mean.size(),
mean.shape))
bn_module.running_mean = torch.from_numpy(mean)
bn_module.running_var = torch.from_numpy(var)
bn_module.weight.data = torch.from_numpy(gamma)
bn_module.bias.data = torch.from_numpy(beta)
def _transfer_residual(prefix, block):
block_module, n_layers = block[0], block[1]
bottleneck = block_module.layers[0]
bottleneck_conv_bn_dic = {prefix + '_1_1x1_reduce': bottleneck.cbr1.cbr_unit,
prefix + '_1_3x3': bottleneck.cbr2.cbr_unit,
prefix + '_1_1x1_proj': bottleneck.cb4.cb_unit,
prefix + '_1_1x1_increase': bottleneck.cb3.cb_unit,}
for k, v in bottleneck_conv_bn_dic.items():
_transfer_conv_bn(k, v)
for layer_idx in range(2, n_layers+1):
residual_layer = block_module.layers[layer_idx-1]
residual_conv_bn_dic = {'_'.join(map(str, [prefix, layer_idx, '1x1_reduce'])): residual_layer.cbr1.cbr_unit,
'_'.join(map(str, [prefix, layer_idx, '3x3'])): residual_layer.cbr2.cbr_unit,
'_'.join(map(str, [prefix, layer_idx, '1x1_increase'])): residual_layer.cb3.cb_unit,}
for k, v in residual_conv_bn_dic.items():
_transfer_conv_bn(k, v)
convbn_layer_mapping = {'conv1_1_3x3_s2': self.convbnrelu1_1.cbr_unit,
'conv1_2_3x3': self.convbnrelu1_2.cbr_unit,
'conv1_3_3x3': self.convbnrelu1_3.cbr_unit,
'conv5_3_pool6_conv': self.pyramid_pooling.paths[0].cbr_unit,
'conv5_3_pool3_conv': self.pyramid_pooling.paths[1].cbr_unit,
'conv5_3_pool2_conv': self.pyramid_pooling.paths[2].cbr_unit,
'conv5_3_pool1_conv': self.pyramid_pooling.paths[3].cbr_unit,
'conv5_4': self.cbr_final.cbr_unit,}
residual_layers = {'conv2': [self.res_block2, self.block_config[0]],
'conv3': [self.res_block3, self.block_config[1]],
'conv4': [self.res_block4, self.block_config[2]],
'conv5': [self.res_block5, self.block_config[3]],}
# Transfer weights for all non-residual conv+bn layers
for k, v in convbn_layer_mapping.items():
_transfer_conv_bn(k, v)
# Transfer weights for final non-bn conv layer
_transfer_conv('conv6', self.classification)
# Transfer weights for all residual layers
for k, v in residual_layers.items():
_transfer_residual(k, v)
def tile_predict(self, img, input_size=[713, 713]):
"""
Predict by takin overlapping tiles from the image.
:param img: np.ndarray with shape [C, H, W] in BGR format
Source: https://github.com/mitmul/chainer-pspnet/blob/master/pspnet.py#L408-L448
Adapted for PyTorch
"""
setattr(self, 'input_size', input_size)
def _pad_img(img):
if img.shape[1] < self.input_size[0]:
pad_h = self.input_size[0] - img.shape[1]
img = np.pad(img, ((0, 0), (0, pad_h), (0, 0)), 'constant')
else:
pad_h = 0
if img.shape[2] < self.input_size[1]:
pad_w = self.input_size[1] - img.shape[2]
img = np.pad(img, ((0, 0), (0, 0), (0, pad_w)), 'constant')
else:
pad_w = 0
return img, pad_h, pad_w
ori_rows, ori_cols = img.shape[1:]
long_size = max(ori_rows, ori_cols)
if long_size > max(self.input_size):
count = np.zeros((ori_rows, ori_cols))
pred = np.zeros((1, self.n_class, ori_rows, ori_cols))
stride_rate = 2 / 3.
stride = (ceil(self.input_size[0] * stride_rate),
ceil(self.input_size[1] * stride_rate))
hh = ceil((ori_rows - self.input_size[0]) / stride[0]) + 1
ww = ceil((ori_cols - self.input_size[1]) / stride[1]) + 1
for yy in range(hh):
for xx in range(ww):
sy, sx = yy * stride[0], xx * stride[1]
ey, ex = sy + self.input_size[0], sx + self.input_size[1]
img_sub = img[:, sy:ey, sx:ex]
img_sub, pad_h, pad_w = _pad_img(img_sub)
inp = Variable(torch.unsqueeze(torch.from_numpy(img_sub), 0).cuda(), volatile=True)
psub1 = F.softmax(self.forward(inp), dim=1).data.cpu().numpy()
psub2 = F.softmax(self.forward(inp[:, :, :, ::-1]), dim=1).data.cpu().numpy()
psub = (psub1 + psub2[:, :, :, ::-1]) / 2.0
if sy + self.input_size[0] > ori_rows:
psub = psub[:, :, :-pad_h, :]
if sx + self.input_size[1] > ori_cols:
psub = psub[:, :, :, :-pad_w]
pred[:, :, sy:ey, sx:ex] = psub
count[sy:ey, sx:ex] += 1
score = (pred / count[None, None, ...]).astype(np.float32)
else:
img, pad_h, pad_w = _pad_img(img)
inp = Variable(torch.unsqueeze(torch.from_numpy(img), 0).cuda(), volatile=True)
pred1 = F.softmax(self.forward(inp), dim=1).data.cpu().numpy()
pred2 = F.softmax(self.forward(inp[:, :, :, ::-1]), dim=1).data.cpu().numpy()
pred = (pred1 + pred2[:, :, :, ::-1]) / 2.0
score = pred[:, :, :self.input_size[0] - pad_h, :self.input_size[1] - pad_w]
score = F.resize_images(score, (ori_rows, ori_cols))[0].data
return score / score.sum(axis=0)
