def conv7x7(in_channels, out_channels, stride=1, padding=0):
"""1x1 convolution"""
weight_shape = (out_channels, in_channels, 7, 7)
weight = Tensor(np.ones(weight_shape).astype(np.float32))
conv = Conv2d(in_channels, out_channels,
kernel_size=7, stride=stride, padding=padding, weight_init=weight, has_bias=False,
pad_mode="same")
conv.conv2d.shard(strategy_weight)
return conv
def conv7x7(in_channels, out_channels, stride=1, padding=0):
"""1x1 convolution"""
weight_shape = (out_channels, in_channels, 7, 7)
weight = variance_scaling_raw(weight_shape)
return Conv2d(in_channels,
out_channels,
kernel_size=7,
stride=stride,
weight_init=weight,
has_bias=False,
pad_mode="same")
def conv3x3(in_channels, out_channels, stride=1, padding=1):
"""3x3 convolution """
weight_shape = (out_channels, in_channels, 3, 3)
weight = variance_scaling_raw(weight_shape)
return Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride,
weight_init=weight,
has_bias=False,
pad_mode="same")
def conv3x3(in_channels, out_channels, stride=1, padding=1):
"""3x3 convolution """
weight_shape = (out_channels, in_channels, 3, 3)
weight = Tensor(np.ones(weight_shape).astype(np.float32))
conv = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=0,
weight_init=weight,
has_bias=False,
pad_mode="same")
conv.conv2d.set_strategy(strategy_weight)
return conv
