def __init__(
self,
in_channels,
channels,
stride=1,
groups=1,
base_width=64,
dilation=1,
norm=M.BatchNorm2d,
):
assert norm is M.BatchNorm2d, "Quant mode only support BatchNorm2d currently."
super(Bottleneck, self).__init__()
width = int(channels * (base_width / 64.0)) * groups
self.conv_bn_relu1 = M.ConvBnRelu2d(in_channels, width, 1, 1, bias=False)
self.conv_bn_relu2 = M.ConvBnRelu2d(
width,
width,
3,
stride,
padding=dilation,
groups=groups,
dilation=dilation,
bias=False,
)
self.conv_bn3 = M.ConvBn2d(width, channels * self.expansion, 1, 1, bias=False)
self.downsample = (
M.Identity()
if in_channels == channels * self.expansion and stride == 1
else M.ConvBn2d(
in_channels, channels * self.expansion, 1, stride, bias=False
)
)
self.add = M.Elemwise("FUSE_ADD_RELU")
def _init_feature_adjust(self):
# feature adjustment
self.r_z_k = M.ConvBn2d(self.channels,
self.feat_channels,
3,
1,
1,
momentum=0.9,
affine=True,
track_running_stats=True)
self.c_z_k = M.ConvBn2d(self.channels,
self.feat_channels,
3,
1,
1,
momentum=0.9,
affine=True,
track_running_stats=True)
self.r_x = M.ConvBn2d(self.channels,
self.feat_channels,
3,
1,
1,
momentum=0.9,
affine=True,
track_running_stats=True)
self.c_x = M.ConvBn2d(self.channels,
self.feat_channels,
3,
1,
1,
momentum=0.9,
affine=True,
track_running_stats=True)
def __init__(
self,
in_channels,
channels,
stride=1,
groups=1,
base_width=64,
dilation=1,
norm=M.BatchNorm2d,
):
assert norm is M.BatchNorm2d, "Quant mode only support BatchNorm2d currently."
super(BasicBlock, self).__init__()
if groups != 1 or base_width != 64:
raise ValueError("BasicBlock only supports groups=1 and base_width=64")
if dilation > 1:
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
self.conv_bn_relu1 = M.ConvBnRelu2d(
in_channels, channels, 3, stride, padding=dilation, bias=False
)
self.conv_bn2 = M.ConvBn2d(channels, channels, 3, 1, padding=1, bias=False)
self.downsample = (
M.Identity()
if in_channels == channels and stride == 1
else M.ConvBn2d(in_channels, channels, 1, stride, bias=False)
)
self.add = M.Elemwise("FUSE_ADD_RELU")
def __init__(self, inp, oup, stride):
super().__init__()
if inp == oup and stride == 1:
self.proj = M.Identity()
else:
self.proj = M.ConvBn2d(inp, oup, 1, stride=stride, bias=False)
self.conv1 = M.ConvBnRelu2d(inp,
oup,
3,
padding=1,
stride=stride,
bias=False)
self.conv2 = M.ConvBn2d(oup, oup, 3, padding=1, stride=1, bias=False)
def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# pw
layers.append(
M.ConvBnRelu2d(inp, hidden_dim, kernel_size=1, bias=False))
layers.extend([
# dw
M.ConvBnRelu2d(
hidden_dim,
hidden_dim,
kernel_size=3,
padding=1,
stride=stride,
groups=hidden_dim,
bias=False,
),
# pw-linear
M.ConvBn2d(hidden_dim, oup, kernel_size=1, bias=False),
])
self.conv = M.Sequential(*layers)
self.add = M.Elemwise("ADD")
def __init__(self):
super().__init__()
self.classifier = None
if dist.get_rank() == 0:
self.features = M.Sequential(
M.ConvBn2d(3, 64, 7, stride=2, padding=3, bias=False),
M.MaxPool2d(kernel_size=3, stride=2, padding=1),
BasicBlock(64, 64, 1),
BasicBlock(64, 64, 1),
)
elif dist.get_rank() == 1:
self.features = M.Sequential(
BasicBlock(64, 128, 2),
BasicBlock(128, 128, 1),
)
elif dist.get_rank() == 2:
self.features = M.Sequential(
BasicBlock(128, 256, 2),
BasicBlock(256, 256, 1),
)
elif dist.get_rank() == 3:
self.features = M.Sequential(
BasicBlock(256, 512, 2),
BasicBlock(512, 512, 1),
)
self.classifier = M.Linear(512, 1000)
def _init_predictor(self):
"""Initialize predictor layers of the head."""
self.conv_cls = M.ConvBn2d(self.feat_channels,
self.cls_out_channels,
kernel_size=1,
stride=1,
padding=0)
self.conv_reg = M.ConvBn2d(self.feat_channels,
2,
kernel_size=1,
stride=1,
padding=0)
self.conv_centerness = M.ConvBn2d(self.feat_channels,
1,
kernel_size=1,
stride=1,
padding=0)
def __init__(self, has_bias=True):
super().__init__()
self.data = np.random.random((1, 3, 224, 224)).astype(np.float32)
self.convbn = M.ConvBn2d(3,
10,
3,
stride=(2, 3),
dilation=(2, 2),
padding=(3, 1),
bias=has_bias)
def __init__(self, ):
super().__init__()
self.conv0 = MyConvBnRelu2d(3, 3, 3, 1, 1)
self.conv1 = M.ConvBn2d(3, 3, 1, 1, 0)
self.conv2 = M.ConvBn2d(3, 3, 1, 1, 0)
self.add = M.Elemwise("FUSE_ADD_RELU")
