def _threshold(raw, input, threshold, value, inplace=False):
# for threshold or relu
if threshold == 0 and value == 0:
x = raw(input, threshold, value, inplace)
bottom_blobs = [log.blobs(input)]
name = log.add_layer(name='relu')
log.add_blobs([x], name='relu_blob')
layer = caffe_net.Layer_param(name=name,
type='ReLU',
bottom=bottom_blobs,
top=[log.blobs(x)])
log.cnet.add_layer(layer)
return x
if value != 0:
raise NotImplemented("value !=0 not implemented in caffe")
x = raw(input, input, threshold, value, inplace)
bottom_blobs = [log.blobs(input)]
layer_name = log.add_layer(name='threshold')
top_blobs = log.add_blobs([x], name='threshold_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Threshold',
bottom=bottom_blobs,
top=top_blobs)
layer.param.threshold_param.threshold = threshold
log.cnet.add_layer(layer)
return x
def _instance_norm(raw,
input,
running_mean=None,
running_var=None,
weight=None,
bias=None,
use_input_stats=True,
momentum=0.1,
eps=1e-5):
# TODO: the batch size!=1 view operations
print(
"WARNING: The Instance Normalization transfers to Caffe using BatchNorm, so the batch size should be 1"
)
if running_var is not None or weight is not None:
# TODO: the affine=True or track_running_stats=True case
raise NotImplementedError(
"not implement the affine=True or track_running_stats=True case InstanceNorm"
)
x = torch.batch_norm(input, weight, bias, running_mean, running_var,
use_input_stats, momentum, eps,
torch.backends.cudnn.enabled)
bottom_blobs = [log.blobs(input)]
layer_name1 = log.add_layer(name='instance_norm')
top_blobs = log.add_blobs([x], name='instance_norm_blob')
layer1 = caffe_net.Layer_param(name=layer_name1,
type='BatchNorm',
bottom=bottom_blobs,
top=top_blobs)
if running_mean is None or running_var is None:
# not use global_stats, normalization is performed over the current mini-batch
layer1.batch_norm_param(use_global_stats=0, eps=eps)
running_mean = torch.zeros(input.size()[1])
running_var = torch.ones(input.size()[1])
else:
layer1.batch_norm_param(use_global_stats=1, eps=eps)
running_mean_clone = running_mean.clone()
running_var_clone = running_var.clone()
layer1.add_data(running_mean_clone.cpu().numpy(),
running_var_clone.cpu().numpy(), np.array([1.0]))
log.cnet.add_layer(layer1)
if weight is not None and bias is not None:
layer_name2 = log.add_layer(name='bn_scale')
layer2 = caffe_net.Layer_param(name=layer_name2,
type='Scale',
bottom=top_blobs,
top=top_blobs)
layer2.param.scale_param.bias_term = True
layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
log.cnet.add_layer(layer2)
return x
def _pool(type, raw, input, x, kernel_size, stride, padding, ceil_mode):
# TODO dilation,ceil_mode,return indices
layer_name = log.add_layer(name='{}_pool'.format(type))
top_blobs = log.add_blobs([x], name='{}_pool_blob'.format(type))
layer = caffe_net.Layer_param(name=layer_name,
type='Pooling',
bottom=[log.blobs(input)],
top=top_blobs)
# TODO w,h different kernel, stride and padding
# processing ceil mode
layer.pool_param(kernel_size=kernel_size,
stride=kernel_size if stride is None else stride,
pad=padding,
type=type.upper())
log.cnet.add_layer(layer)
if ceil_mode == False and stride is not None:
oheight = (input.size()[2] - _pair(kernel_size)[0] +
2 * _pair(padding)[0]) % (_pair(stride)[0])
owidth = (input.size()[3] - _pair(kernel_size)[1] +
2 * _pair(padding)[1]) % (_pair(stride)[1])
if oheight != 0 or owidth != 0:
caffe_out = raw(input,
kernel_size,
stride,
padding,
ceil_mode=True)
print(
"WARNING: the output shape miss match at {}: "
"input {} output---Pytorch:{}---Caffe:{}\n"
"This is caused by the different implementation that ceil mode in caffe and the floor mode in pytorch.\n"
"You can add the clip layer in caffe prototxt manually if shape mismatch error is caused in caffe. "
.format(layer_name, input.size(), x.size(), caffe_out.size()))
def _interpolate(raw,
input,
size=None,
scale_factor=None,
mode='nearest',
align_corners=None):
# 定义的参数包括 scale,即输出与输入的尺寸比例,如 2;scale_h、scale_w,
# 同 scale,分别为 h、w 方向上的尺寸比例;pad_out_h、pad_out_w,仅在 scale 为 2 时
# 有用,对输出进行额外 padding 在 h、w 方向上的数值;upsample_h、upsample_w,输
# 出图像尺寸的数值。在 Upsample 的相关代码中,推荐仅仅使用 upsample_h、
# upsample_w 准确定义 Upsample 层的输出尺寸,其他所有的参数都不推荐继续使用。
# for nearest _interpolate
if mode != "nearest" or align_corners != None:
raise NotImplementedError("not implement F.interpolate totoaly")
x = raw(input, size, scale_factor, mode)
layer_name = log.add_layer(name='upsample')
top_blobs = log.add_blobs([x], name='upsample_blob'.format(type))
layer = caffe_net.Layer_param(name=layer_name,
type='Upsample',
bottom=[log.blobs(input)],
top=top_blobs)
layer.upsample_param(size=(input.size(2), input.size(3)),
scale_factor=scale_factor)
log.cnet.add_layer(layer)
return x
def _conv_transpose2d(raw,
input,
weight,
bias=None,
stride=1,
padding=0,
output_padding=0,
groups=1,
dilation=1):
x = raw(input, weight, bias, stride, padding, output_padding, groups,
dilation)
name = log.add_layer(name='conv_transpose')
log.add_blobs([x], name='conv_transpose_blob')
layer = caffe_net.Layer_param(name=name,
type='Deconvolution',
bottom=[log.blobs(input)],
top=[log.blobs(x)])
layer.conv_param(x.size()[1],
weight.size()[2:],
stride=_pair(stride),
pad=_pair(padding),
dilation=_pair(dilation),
bias_term=bias is not None)
if bias is not None:
layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
else:
layer.param.convolution_param.bias_term = False
layer.add_data(weight.cpu().data.numpy())
log.cnet.add_layer(layer)
return x
def _leaky_relu(raw, input, negative_slope=0.01, inplace=False):
x = raw(input, negative_slope)
name = log.add_layer(name='leaky_relu')
log.add_blobs([x], name='leaky_relu_blob')
layer = caffe_net.Layer_param(name=name,
type='ReLU',
bottom=[log.blobs(input)],
top=[log.blobs(x)])
layer.param.relu_param.negative_slope = negative_slope
log.cnet.add_layer(layer)
return x
def _tanh(raw, input):
# for tanh activation
x = raw(input)
name = log.add_layer(name='tanh')
log.add_blobs([x], name='tanh_blob')
layer = caffe_net.Layer_param(name=name,
type='TanH',
bottom=[log.blobs(input)],
top=[log.blobs(x)])
log.cnet.add_layer(layer)
return x
def _relu(raw, input, inplace=False):
# for threshold or prelu
x = raw(input, False)
name = log.add_layer(name='relu')
log.add_blobs([x], name='relu_blob')
layer = caffe_net.Layer_param(name=name,
type='ReLU',
bottom=[log.blobs(input)],
top=[log.blobs(x)])
log.cnet.add_layer(layer)
return x
def _batch_norm(raw,
input,
running_mean,
running_var,
weight=None,
bias=None,
training=False,
momentum=0.1,
eps=1e-5):
# because the runing_mean and runing_var will be changed after the _batch_norm operation, we first save the parameters
x = raw(input, running_mean, running_var, weight, bias, training, momentum,
eps)
bottom_blobs = [log.blobs(input)]
layer_name1 = log.add_layer(name='batch_norm')
top_blobs = log.add_blobs([x], name='batch_norm_blob')
layer1 = caffe_net.Layer_param(name=layer_name1,
type='BatchNorm',
bottom=bottom_blobs,
top=top_blobs)
if running_mean is None or running_var is None:
# not use global_stats, normalization is performed over the current mini-batch
layer1.batch_norm_param(use_global_stats=0, eps=eps)
else:
layer1.batch_norm_param(use_global_stats=1, eps=eps)
running_mean_clone = running_mean.clone()
running_var_clone = running_var.clone()
layer1.add_data(running_mean_clone.cpu().numpy(),
running_var_clone.cpu().numpy(), np.array([1.0]))
log.cnet.add_layer(layer1)
if weight is not None and bias is not None:
layer_name2 = log.add_layer(name='bn_scale')
layer2 = caffe_net.Layer_param(name=layer_name2,
type='Scale',
bottom=top_blobs,
top=top_blobs)
layer2.param.scale_param.bias_term = True
layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
log.cnet.add_layer(layer2)
return x
def _mul(input, *args):
x = raw__sub__(input, *args)
if not NET_INITTED:
return x
layer_name = log.add_layer(name='mul')
top_blobs = log.add_blobs([x], name='mul_blob')
layer = caffe_net.Layer_param(
name=layer_name,
type='Eltwise',
bottom=[log.blobs(input), log.blobs(args[0])],
top=top_blobs)
layer.param.eltwise_param.operation = 0 # product is 1
log.cnet.add_layer(layer)
return x
def _cat(raw, inputs, dimension=0):
x = raw(inputs, dimension)
bottom_blobs = []
for input in inputs:
bottom_blobs.append(log.blobs(input))
layer_name = log.add_layer(name='cat')
top_blobs = log.add_blobs([x], name='cat_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Concat',
bottom=bottom_blobs,
top=top_blobs)
layer.param.concat_param.axis = dimension
log.cnet.add_layer(layer)
return x
def _dropout(raw, input, p=0.5, training=False, inplace=False):
x = raw(input, p, training, inplace)
bottom_blobs = [log.blobs(input)]
layer_name = log.add_layer(name='dropout')
top_blobs = log.add_blobs([x], name=bottom_blobs[0], with_num=False)
layer = caffe_net.Layer_param(name=layer_name,
type='Dropout',
bottom=bottom_blobs,
top=top_blobs)
layer.param.dropout_param.dropout_ratio = p
layer.param.include.extend([caffe_net.pb.NetStateRule(phase=0)
]) # 1 for test, 0 for train
log.cnet.add_layer(layer)
return x
def _isub(input, *args):
x = raw__isub__(input, *args)
if not NET_INITTED:
return x
x = x.clone()
layer_name = log.add_layer(name='sub')
top_blobs = log.add_blobs([x], name='sub_blob')
layer = caffe_net.Layer_param(
name=layer_name,
type='Eltwise',
bottom=[log.blobs(input), log.blobs(args[0])],
top=top_blobs)
layer.param.eltwise_param.operation = 1 # sum is 1
log.cnet.add_layer(layer)
return x
def _split(raw, tensor, split_size, dim=0):
# split in pytorch is slice in caffe
x = raw(tensor, split_size, dim)
layer_name = log.add_layer('split')
top_blobs = log.add_blobs(x, name='split_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Slice',
bottom=[log.blobs(tensor)],
top=top_blobs)
slice_num = int(np.floor(tensor.size()[dim] / split_size))
slice_param = caffe_net.pb.SliceParameter(
axis=dim, slice_point=[split_size * i for i in range(1, slice_num)])
layer.param.slice_param.CopyFrom(slice_param)
log.cnet.add_layer(layer)
return x
def _softmax(raw, input, dim=None, _stacklevel=3):
# for F.softmax
x = raw(input, dim=dim)
if dim is None:
dim = F._get_softmax_dim('softmax', input.dim(), _stacklevel)
bottom_blobs = [log.blobs(input)]
name = log.add_layer(name='softmax')
log.add_blobs([x], name='softmax_blob')
layer = caffe_net.Layer_param(name=name,
type='Softmax',
bottom=bottom_blobs,
top=[log.blobs(x)])
layer.param.softmax_param.axis = dim
log.cnet.add_layer(layer)
return x
def _linear(raw, input, weight, bias=None):
x = raw(input, weight, bias)
layer_name = log.add_layer(name='fc')
top_blobs = log.add_blobs([x], name='fc_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='InnerProduct',
bottom=[log.blobs(input)],
top=top_blobs)
layer.fc_param(x.size()[1], has_bias=bias is not None)
if bias is not None:
layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy())
else:
layer.add_data(weight.cpu().data.numpy())
log.cnet.add_layer(layer)
return x
def _view(input, *args):
ipdb.set_trace()
x = raw_view(input, *args)
if not NET_INITTED:
return x
layer_name = log.add_layer(name='view')
top_blobs = log.add_blobs([x], name='view_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Reshape',
bottom=[log.blobs(input)],
top=top_blobs)
# TODO: reshpae added to nn_tools layer
dims = list(args)
dims[0] = 0 # the first dim should be batch_size
layer.param.reshape_param.shape.CopyFrom(caffe_net.pb.BlobShape(dim=dims))
log.cnet.add_layer(layer)
return x
def _prelu(raw, input, weight):
# for threshold or prelu
x = raw(input, weight)
bottom_blobs = [log.blobs(input)]
name = log.add_layer(name='prelu')
log.add_blobs([x], name='prelu_blob')
layer = caffe_net.Layer_param(name=name,
type='PReLU',
bottom=bottom_blobs,
top=[log.blobs(x)])
if weight.size()[0] == 1:
layer.param.prelu_param.channel_shared = True
layer.add_data(weight.cpu().data.numpy()[0])
else:
layer.add_data(weight.cpu().data.numpy())
log.cnet.add_layer(layer)
return x
def _max(raw, *args):
x = raw(*args)
if len(args) == 1:
# TODO max in one tensor
assert NotImplementedError
else:
bottom_blobs = []
for arg in args:
bottom_blobs.append(log.blobs(arg))
layer_name = log.add_layer(name='max')
top_blobs = log.add_blobs([x], name='max_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Eltwise',
bottom=bottom_blobs,
top=top_blobs)
layer.param.eltwise_param.operation = 2
log.cnet.add_layer(layer)
return x
def _mean(input, *args, **kwargs):
x = raw_mean(input, *args, **kwargs)
if not NET_INITTED:
return x
layer_name = log.add_layer(name='mean')
top_blobs = log.add_blobs([x], name='mean_blob')
layer = caffe_net.Layer_param(name=layer_name,
type='Reduction',
bottom=[log.blobs(input)],
top=top_blobs)
if len(args) == 1:
dim = args[0]
elif 'dim' in kwargs:
dim = kwargs['dim']
else:
raise NotImplementedError('mean operation must specify a dim')
layer.param.reduction_param.operation = 4
layer.param.reduction_param.axis = dim
log.cnet.add_layer(layer)
return x
