Hi @ptrblck , the model.apply(...) is used to set the ‘width_mult’ attribute for different widths. It is related to my custom convolution. I define my custom convolution as follows
class USConv2d(nn.Conv2d):
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1, depthwise=False, bias=True,
us=[True, True], ratio=[1, 1]):
in_channels_max = in_channels
out_channels_max = out_channels
if us[0]:
in_channels_max = int(make_divisible(
in_channels
* FLAGS.width_mult
/ ratio[0]) * ratio[0])
if us[1]:
out_channels_max = int(make_divisible(
out_channels
* FLAGS.width_mult
/ ratio[1]) * ratio[1])
groups = in_channels_max if depthwise else 1
super(USConv2d, self).__init__(
in_channels_max, out_channels_max,
kernel_size, stride=stride, padding=padding, dilation=dilation,
groups=groups, bias=bias)
self.depthwise = depthwise
self.in_channels_basic = in_channels
self.out_channels_basic = out_channels
self.width_mult = None
self.us = us
self.ratio = ratio
def forward(self, input):
in_channels = self.in_channels_basic
out_channels = self.out_channels_basic
if self.us[0]:
in_channels = int(make_divisible(
self.in_channels_basic
* self.width_mult
/ self.ratio[0]) * self.ratio[0])
if self.us[1]:
out_channels = int(make_divisible(
self.out_channels_basic
* self.width_mult
/ self.ratio[1]) * self.ratio[1])
self.groups = in_channels if self.depthwise else 1
if self.bias is not None:
y = nn.functional.conv2d(
input, self.weight[:out_channels, :in_channels, :, :], self.bias[:out_channels], self.stride, self.padding,
self.dilation, self.groups)
else:
y = nn.functional.conv2d(
input, self.weight[:out_channels, :in_channels, :, :], self.bias, self.stride,
self.padding, self.dilation, self.groups)
if getattr(FLAGS, 'conv_averaged', False):
y = y * (max(self.in_channels_list)/self.in_channels)
return y
The width_mult is to calculate the number of channels used used for current width. Similarly, I define the custom BN and Linear layer. Codes are modified from the slimmable networks
