def __init__(self,
in_channels,
out_channels,
stride,
dropout,
use_se=False):
super().__init__()
self.use_se = use_se
self.norm1 = Norm(in_channels)
self.act1 = Act()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride)
self.norm2 = Norm(out_channels)
self.act2 = Act()
self.dropout = Dropout(dropout) if dropout else Identity()
self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
if self.use_se:
self.se = SELayer(out_channels, reduction=8)
if stride != 1:
assert in_channels != out_channels
self.shortcut = Sequential([
Pool2d(2, 2, type='avg'),
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'),
])
else:
self.shortcut = Identity()
def __init__(self, in_channels, channels, stride, channels_per_group, cardinality,
start_block=False, end_block=False, exclude_bn0=False):
super().__init__()
out_channels = channels * self.expansion
width = channels_per_group * cardinality
if not start_block and not exclude_bn0:
self.bn0 = Norm(in_channels)
if not start_block:
self.act0 = Act()
self.conv1 = Conv2d(in_channels, width, kernel_size=1)
self.bn1 = Norm(width)
self.act1 = Act()
self.conv2 = Conv2d(width, width, kernel_size=3, stride=stride, groups=cardinality,
norm='def', act='def')
self.conv3 = Conv2d(width, out_channels, kernel_size=1)
if start_block:
self.bn3 = Norm(out_channels)
if end_block:
self.bn3 = Norm(out_channels)
self.act3 = Act()
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(3, 2, type='max'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.start_block = start_block
self.end_block = end_block
self.exclude_bn0 = exclude_bn0
def __init__(self, in_channels, out_channels, dropout):
layers = [
Norm(in_channels),
Act(),
Conv2d(in_channels, out_channels, kernel_size=3),
Norm(out_channels),
Act(),
Conv2d(out_channels, out_channels, kernel_size=3),
]
if dropout:
layers.insert(5, Dropout(dropout))
super().__init__(layers)
def __init__(self,
in_channels,
channels,
stride,
dropout=0,
drop_path=0,
avd=False,
start_block=False,
end_block=False,
exclude_bn0=False):
super().__init__()
out_channels = channels * self.expansion
if not start_block and not exclude_bn0:
self.bn0 = Norm(in_channels)
if not start_block:
self.act0 = Act()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride)
self.bn1 = Norm(out_channels)
self.act1 = Act()
self.dropout = Dropout(dropout) if dropout else Identity()
self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
if start_block:
self.bn2 = Norm(out_channels)
self.drop_path = DropPath(drop_path) if drop_path else Identity()
if end_block:
self.bn2 = Norm(out_channels)
self.act2 = Act()
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(3, 2, type='max'))
if in_channels != out_channels:
shortcut.append(
Conv2d(in_channels,
out_channels,
kernel_size=1,
norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.start_block = start_block
self.end_block = end_block
self.exclude_bn0 = exclude_bn0
def __init__(self, in_channels, out_channels, dropout, reduction):
layers = [
Norm(in_channels),
Act(),
Conv2d(in_channels, out_channels, kernel_size=3),
Norm(out_channels),
Act(),
Conv2d(out_channels, out_channels, kernel_size=3),
]
if dropout:
layers.insert(5, Dropout(dropout))
layers.append(SELayer(out_channels, reduction=reduction))
super().__init__(layers)
def __init__(self, d_model, num_heads, dff, drop_rate, attn_drop_rate, activation='gelu', **kwargs):
super().__init__(**kwargs)
self.d_model = d_model
self.num_heads = num_heads
self.dff = dff
self.drop_rate = drop_rate
self.attn_drop_rate = attn_drop_rate
self.activation = activation
self.mha = MultiHeadAttention(d_model, num_heads, attn_drop_rate)
self.ln1 = Norm(type='ln')
self.ffn = FFN(d_model, dff, drop_rate, activation)
self.ln2 = Norm(type='ln')
def __init__(self,
in_channels,
channels,
stride,
base_width,
scale,
cardinality,
end_block=False,
exclude_bn0=False):
super().__init__()
out_channels = channels * self.expansion
start_block = stride != 1 or in_channels != out_channels
width = math.floor(channels * (base_width / 64)) * scale * cardinality
if not start_block and not exclude_bn0:
self.bn0 = Norm(in_channels)
if not start_block:
self.act0 = Act()
self.conv1 = Conv2d(in_channels, width, kernel_size=1)
self.bn1 = Norm(width)
self.act1 = Act()
self.conv2 = Res2Conv(width,
width,
kernel_size=3,
stride=stride,
scale=scale,
groups=cardinality,
norm='def',
act='def',
start_block=start_block)
self.conv3 = Conv2d(width, out_channels, kernel_size=1)
if start_block:
self.bn3 = Norm(out_channels)
if end_block:
self.bn3 = Norm(out_channels)
self.act3 = Act()
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.start_block = start_block
self.end_block = end_block
self.exclude_bn0 = exclude_bn0
def __init__(self, C, layers, steps=3, stem_multiplier=3, drop_path=0.6, num_classes=10):
super().__init__()
self._C = C
self._steps = steps
self._drop_path = drop_path
C_curr = stem_multiplier * C
self.stem = Sequential([
Conv2d(3, C_curr, 3, bias=False),
Norm(3, 'def', affine=True),
])
C_prev, C_curr = C_curr, C
self.cells = []
for i in range(layers):
if i in [layers // 3, 2 * layers // 3]:
C_curr *= 2
cell = BasicBlock(C_prev, C_curr, stride=2)
else:
cell = Cell(steps, C_prev, C_curr, drop_path)
self.cells.append(cell)
C_prev = C_curr
self.global_pool = GlobalAvgPool()
self.fc = Linear(C_prev, num_classes)
k = sum(1 + i for i in range(self._steps))
num_ops = len(get_primitives())
self.alphas_normal = self.add_weight(
'alphas_normal', (k, num_ops), initializer=RandomNormal(stddev=1e-2), trainable=True,
)
def __init__(self, in_channels, out_channels, stride, groups, reduction,
zero_init_residual):
super().__init__()
se_channels = in_channels // reduction
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=1,
norm='def',
act='def')
self.conv2 = Conv2d(out_channels,
out_channels,
kernel_size=3,
stride=stride,
groups=groups,
norm='def',
act='def')
self.se = SELayer(out_channels, se_channels)
self.conv3 = Sequential([
Conv2d(out_channels, out_channels, kernel_size=1, bias=False),
Norm(out_channels,
gamma_init='zeros' if zero_init_residual else 'ones')
])
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self, C_in, C_out, kernel_size, stride):
super().__init__([
Act(),
Conv2d(C_in,
C_in,
kernel_size,
stride=stride,
groups=C_in,
bias=False),
Conv2d(C_in, C_in, 1, bias=False),
Norm(C_in),
Act(),
Conv2d(C_in, C_in, kernel_size, 1, groups=C_in, bias=False),
Conv2d(C_in, C_out, 1, bias=False),
Norm(C_out),
])
def __init__(self, start_channels, widening_fractor, depth, groups, radix, drop_path, num_classes=10):
super().__init__()
num_layers = [(depth - 2) // 9] * 3
strides = [1, 2, 2]
self.add_channel = widening_fractor / sum(num_layers)
self.in_channels = start_channels
self.channels = start_channels
layers = [Conv2d(3, start_channels, kernel_size=3, norm='def')]
for i, (n, s) in enumerate(zip(num_layers, strides)):
layers.append(self._make_layer(n, groups, stride=s, radix=radix, drop_path=drop_path))
layers.append(Sequential([
Norm(self.in_channels),
Act(),
]))
self.features = Sequential(layers)
assert (start_channels + widening_fractor) * Bottleneck.expansion == self.in_channels
self.final_pool = GlobalAvgPool()
self.fc = Linear(self.in_channels, num_classes)
def __init__(self,
in_channels,
channels,
groups,
stride=1,
se_reduction=4,
drop_path=0.2):
super().__init__()
branch1 = [
Norm(in_channels),
Conv2d(in_channels,
channels,
kernel_size=1,
norm='default',
act='default'),
*([Pool2d(3, 2)] if stride != 1 else []),
Conv2d(channels,
channels,
kernel_size=3,
groups=groups,
norm='default',
act='default'),
*([SELayer(channels, se_reduction, groups)]
if se_reduction else []),
Conv2d(channels, channels, kernel_size=1, norm='default'),
*([DropPath(drop_path)] if drop_path and stride == 1 else []),
]
self.branch1 = Sequential(branch1)
self.branch2 = Shortcut(in_channels, channels, stride)
def __init__(self, in_channels, out_channels, dropout, use_se, drop_path):
layers = [
Norm(in_channels),
Act(),
Conv2d(in_channels, out_channels, kernel_size=3),
Norm(out_channels),
Act(),
Conv2d(out_channels, out_channels, kernel_size=3),
]
if dropout:
layers.insert(5, Dropout(dropout))
if use_se:
layers.append(SELayer(out_channels, reduction=8))
if drop_path:
layers.append(DropPath(drop_path))
super().__init__(layers)
def __init__(self, depth, k, dropout=0, reduction=8, num_classes=10):
super().__init__()
num_blocks = (depth - 4) // 6
self.conv = Conv2d(3, self.stages[0], kernel_size=3)
self.layer1 = self._make_layer(self.stages[0] * 1,
self.stages[1] * k,
num_blocks,
stride=1,
dropout=dropout,
reduction=reduction)
self.layer2 = self._make_layer(self.stages[1] * k,
self.stages[2] * k,
num_blocks,
stride=2,
dropout=dropout,
reduction=reduction)
self.layer3 = self._make_layer(self.stages[2] * k,
self.stages[3] * k,
num_blocks,
stride=2,
dropout=dropout,
reduction=reduction)
self.norm = Norm(self.stages[3] * k)
self.act = Act()
self.avgpool = GlobalAvgPool()
self.fc = Linear(self.stages[3] * k, num_classes)
def __init__(self, in_channels, out_channels, stride, groups, use_se):
super().__init__()
self.use_se = use_se
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=1,
norm='def',
act='def')
self.conv2 = Conv2d(out_channels,
out_channels,
kernel_size=3,
stride=stride,
groups=groups,
norm='def',
act='def')
if self.use_se:
self.se = SELayer(out_channels, 4)
self.conv3 = Sequential(
Conv2d(out_channels, out_channels, kernel_size=1, bias=False),
Norm(out_channels, gamma_init='zeros'))
if stride != 1 or in_channels != out_channels:
self.shortcut = Conv2d(in_channels,
out_channels,
kernel_size=1,
norm='def')
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self, in_channels, out_channels, stride, groups):
super().__init__()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=1,
norm='def',
act='def')
self.conv2 = Conv2d(out_channels,
out_channels,
kernel_size=3,
stride=stride,
groups=groups,
norm='def',
act='def')
self.conv3 = Sequential([
Conv2d(out_channels, out_channels, kernel_size=1, bias=False),
Norm(out_channels)
])
self.eca = ECALayer()
if stride != 1 or in_channels != out_channels:
self.shortcut = Conv2d(in_channels,
out_channels,
stride=stride,
kernel_size=1,
norm='def')
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self,
in_channels,
channels,
stride,
zero_init_residual=True,
reduction=16):
super().__init__()
out_channels = channels * self.expansion
self.conv1 = Conv2d(in_channels,
channels,
kernel_size=1,
norm='def',
act='def')
self.conv2 = Conv2d(channels,
channels,
kernel_size=3,
stride=stride,
norm='def',
act='def')
self.conv3 = Conv2d(channels, out_channels, kernel_size=1)
self.bn3 = Norm(out_channels,
gamma_init='zeros' if zero_init_residual else 'ones')
self.se = SELayer(out_channels, reduction=reduction)
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self, in_channels, channels, stride, dilation=1, base_width=26, scale=4,
zero_init_residual=True, avd=False):
super().__init__()
out_channels = channels * self.expansion
start_block = stride != 1 or in_channels != out_channels
width = math.floor(channels * (base_width / 64)) * scale
self.conv1 = Conv2d(in_channels, width, kernel_size=1,
norm='def', act='def')
if avd and stride != 1:
self.conv2 = Sequential([
Pool2d(3, stride=stride, type='avg'),
Res2Conv(width, width, kernel_size=3, stride=1, dilation=dilation, scale=scale,
groups=1, start_block=start_block, norm='def', act='def'),
])
else:
self.conv2 = Res2Conv(width, width, kernel_size=3, stride=stride, dilation=dilation,
scale=scale, groups=1, start_block=start_block, norm='def', act='def')
self.conv3 = Conv2d(width, out_channels, kernel_size=1)
self.bn3 = Norm(out_channels, gamma_init='zeros' if zero_init_residual else 'ones')
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self, in_channels, channels, stride, base_width, splits, zero_init_residual, genotype):
super().__init__()
self.stride = stride
out_channels = channels * self.expansion
width = math.floor(out_channels // self.expansion * (base_width / 64)) * splits
self.conv1 = Conv2d(in_channels, width, kernel_size=1,
norm='def', act='def')
if stride == 1:
self.conv2 = PPConv(width, splits=splits, genotype=genotype)
else:
self.conv2 = Conv2d(width, width, kernel_size=3, stride=2, groups=splits,
norm='def', act='def')
self.conv3 = Conv2d(width, out_channels, kernel_size=1)
self.bn3 = Norm(out_channels, gamma_init='zeros' if zero_init_residual else 'ones')
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.act = Act()
def __init__(self, in_channels, out_channels, stride, dropout):
super().__init__()
self.norm1 = Norm(in_channels)
self.act1 = Act()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride)
self.norm2 = Norm(out_channels)
self.act2 = Act()
self.dropout = Dropout(dropout) if dropout else Identity()
self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
self.shortcut = Conv2d(in_channels,
out_channels,
kernel_size=1,
stride=stride)
def __init__(self, C, stride, drop_path):
super().__init__()
self.stride = stride
self._ops = []
for i, primitive in enumerate(get_primitives()):
if drop_path:
op = Sequential([
OPS[primitive](C, stride),
])
if 'pool' in primitive:
op.add(Norm(C))
op.add(DropPath(drop_path))
else:
op = OPS[primitive](C, stride)
if 'pool' in primitive:
op = Sequential([op, Norm(C)])
self._ops.append(op)
def __init__(self,
in_channels,
channels,
stride,
start_block=False,
end_block=False,
exclude_bn0=False,
conv_cls=Conv2d):
super().__init__()
out_channels = channels * self.expansion
width = getattr(self, "width", channels)
if not start_block and not exclude_bn0:
self.bn0 = Norm(in_channels)
if not start_block:
self.act0 = Act()
self.conv1 = Conv2d(in_channels, width, kernel_size=1)
self.bn1 = Norm(width)
self.act1 = Act()
self.conv2 = conv_cls(in_channels=width,
out_channels=width,
kernel_size=3,
stride=stride,
norm='def',
act='def',
start_block=start_block)
self.conv3 = Conv2d(width, out_channels, kernel_size=1)
if start_block:
self.bn3 = Norm(out_channels)
if end_block:
self.bn3 = Norm(out_channels)
self.act3 = Act()
if stride != 1 or in_channels != out_channels:
shortcut = []
if stride != 1:
shortcut.append(Pool2d(2, 2, type='avg'))
shortcut.append(
Conv2d(in_channels, out_channels, kernel_size=1, norm='def'))
self.shortcut = Sequential(shortcut)
else:
self.shortcut = Identity()
self.start_block = start_block
self.end_block = end_block
self.exclude_bn0 = exclude_bn0
def __init__(self,
in_channels,
out_channels,
stride,
askc_type='DirectAdd'):
super().__init__()
self.norm1 = Norm(in_channels)
self.act1 = Act()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride,
bias=False)
self.norm2 = Norm(out_channels)
self.act2 = Act()
self.conv2 = Conv2d(out_channels,
out_channels,
kernel_size=3,
bias=False)
if in_channels != out_channels or stride == 2:
if stride == 2:
self.shortcut = Sequential([
Pool2d(2, 2, type='avg'),
Conv2d(in_channels,
out_channels,
kernel_size=1,
bias=False),
])
else:
self.shortcut = Conv2d(in_channels,
out_channels,
kernel_size=1,
bias=False)
else:
self.shortcut = Identity()
if askc_type == 'DirectAdd':
self.attention = DirectAddFuse()
elif askc_type == 'iAFF':
self.attention = iAFF(out_channels)
elif askc_type == 'AFF':
self.attention = AFF(out_channels)
else:
raise ValueError('Unknown askc_type')
def __init__(self, C_in, C_out):
super().__init__()
assert C_out % 2 == 0
self.act = Act()
self.slice = Slice([1, 1, 0], [-1, -1, -1])
self.conv1 = Conv2d(C_in, C_out // 2, 1, stride=2, bias=False)
self.conv2 = Conv2d(C_in, C_out // 2, 1, stride=2, bias=False)
self.norm = Norm(C_out)
self.concat = Concatenate()
def __init__(self, in_channels, num_classes):
layers = [
Norm(in_channels),
Pool2d(5, 3, padding=0, type='avg'),
Conv2d(in_channels, 128, 1, norm='def', act='def'),
Conv2d(128, 768, 2, padding=0, norm='def', act='def'),
Flatten(),
Linear(768, num_classes),
]
super().__init__(layers)
def __init__(self, in_channels, out_channels, stride):
super().__init__()
self.norm1 = Norm(in_channels)
self.act1 = Act()
self.conv1 = Conv2d(in_channels,
out_channels,
kernel_size=3,
stride=stride)
self.norm2 = Norm(out_channels)
self.act2 = Act()
self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3)
if stride != 1 or in_channels != out_channels:
self.shortcut = Conv2d(in_channels,
out_channels,
kernel_size=1,
stride=stride)
self.rezero = ReZero()
def __init__(self,
C,
layers,
steps=4,
multiplier=4,
stem_multiplier=3,
drop_path=0,
num_classes=10):
super().__init__()
self._C = C
self._steps = steps
self._multiplier = multiplier
self._drop_path = drop_path
C_curr = stem_multiplier * C
self.stem = Sequential([
Conv2d(3, C_curr, 3, bias=False),
Norm(C_curr, 'def', affine=True),
])
C_prev_prev, C_prev, C_curr = C_curr, C_curr, C
self.cells = []
reduction_prev = False
for i in range(layers):
if i in [layers // 3, 2 * layers // 3]:
C_curr *= 2
reduction = True
else:
reduction = False
cell = Cell(steps, multiplier, C_prev_prev, C_prev, C_curr,
reduction, reduction_prev, drop_path)
reduction_prev = reduction
self.cells.append(cell)
C_prev_prev, C_prev = C_prev, multiplier * C_curr
self.avg_pool = GlobalAvgPool()
self.classifier = Linear(C_prev, num_classes)
k = sum(2 + i for i in range(self._steps))
num_ops = len(get_primitives())
self.alphas_normal = self.add_weight(
'alphas_normal', (k, num_ops),
initializer=RandomNormal(stddev=1e-2),
trainable=True,
experimental_autocast=False)
self.alphas_reduce = self.add_weight(
'alphas_reduce', (k, num_ops),
initializer=RandomNormal(stddev=1e-2),
trainable=True,
experimental_autocast=False)
self.tau = self.add_weight('tau', (),
initializer=Constant(1.0),
trainable=False,
experimental_autocast=False)
def __init__(self,
start_channels,
alpha,
depth,
block='bottleneck',
p_shakedrop=0.5,
num_classes=10):
super().__init__()
if block == 'basic':
num_layers = [(depth - 2) // 6] * 3
block = BasicBlock
elif block == 'bottleneck':
num_layers = [(depth - 2) // 9] * 3
block = Bottleneck
else:
raise ValueError("block must be `basic` or `bottleneck`, got %s" %
block)
self.num_layers = num_layers
strides = [1, 2, 2]
add_channel = alpha / sum(num_layers)
in_channels = start_channels
self.init_block = Conv2d(3, start_channels, kernel_size=3, norm='def')
channels = start_channels
k = 1
units = []
for n, s in zip(num_layers, strides):
for i in range(n):
stride = s if i == 0 else 1
channels = channels + add_channel
units.append(
block(in_channels,
rd(channels),
stride=stride,
p_shakedrop=k / sum(num_layers) * p_shakedrop))
in_channels = rd(channels) * block.expansion
k += 1
self.units = units
self.post_activ = Sequential([
Norm(in_channels),
Act(),
])
assert (start_channels + alpha) * block.expansion == in_channels
self.final_pool = GlobalAvgPool()
self.fc = Linear(in_channels, num_classes)
def __init__(self, C, stride, k):
super().__init__()
self.stride = stride
self.mp = Pool2d(2, 2, type='max')
self.k = k
self._ops = []
self._channels = C // k
for i, primitive in enumerate(get_primitives()):
op = OPS[primitive](self._channels, stride)
if 'pool' in primitive:
op = Sequential([op, Norm(self._channels)])
self._ops.append(op)
def __init__(self, C_prev_prev, C_prev, C):
super().__init__()
self.reduction = True
self.preprocess0 = ReLUConvBN(C_prev_prev, C, 1)
self.preprocess1 = ReLUConvBN(C_prev, C, 1)
self.branch_a1 = Sequential([
Act(),
Conv2d(C, C, (1, 3), stride=(1, 2), groups=8, bias=False),
Conv2d(C, C, (3, 1), stride=(2, 1), groups=8, bias=False),
Norm(C, affine=True),
Act(),
Conv2d(C, C, 1),
Norm(C, affine=True),
])
self.branch_a2 = Sequential([
Pool2d(3, stride=2, type='max'),
Norm(C, affine=True),
])
self.branch_b1 = Sequential([
Act(),
Conv2d(C, C, (1, 3), stride=(1, 2), groups=8, bias=False),
Conv2d(C, C, (3, 1), stride=(2, 1), groups=8, bias=False),
Norm(C, affine=True),
Act(),
Conv2d(C, C, 1),
Norm(C, affine=True),
])
self.branch_b2 = Sequential([
Pool2d(3, stride=2, type='max'),
Norm(C, affine=True),
])
