def sepconv2d(cin, cout=None, ksize=3, stride=1, padding=None, affine=True):
if cout is None: cout = cin
if padding is None: padding = ksize // 2
layer = nn.Sequential(
nn.ReLU(inplace=False),
init_default(
nn.Conv2d(cin,
cin,
ksize,
stride=stride,
padding=padding,
groups=cin,
bias=False), nn.init.kaiming_normal_),
init_default(nn.Conv2d(cin, cin, 1, padding=0, bias=False),
nn.init.kaiming_normal_),
nn.BatchNorm2d(cin, affine=affine), nn.ReLU(inplace=False),
init_default(
nn.Conv2d(cin,
cin,
ksize,
stride=1,
padding=padding,
groups=cin,
bias=False), nn.init.kaiming_normal_),
init_default(nn.Conv2d(cin, cout, 1, padding=0, bias=False),
nn.init.kaiming_normal_),
nn.BatchNorm2d(cout, affine=affine))
return layer
def conv2dpool(cin, cout, pool_type, bn=NormType.Batch):
assert pool_type in ['avg', 'max']
if pool_type == 'max':
return nn.Sequential(
nn.MaxPool2d(2, stride=2),
init_default(nn.Conv2d(cin, cout, 1, bias=False),
nn.init.kaiming_normal_),
batchnorm_2d(cout, norm_type=bn))
if pool_type == 'avg':
return nn.Sequential(
nn.AvgPool2d(2, stride=2, ceil_mode=True, count_include_pad=False),
init_default(nn.Conv2d(cin, cout, 1, bias=False),
nn.init.kaiming_normal_),
batchnorm_2d(cout, norm_type=bn))
def conv2d(cin,
cout=None,
ksize=3,
stride=1,
padding=None,
dilation=None,
groups=None,
use_relu=True,
use_bn=True,
bn=NormType.Batch,
bias=False):
if cout is None: cout = cin
if padding is None: padding = ksize // 2
if dilation is None: dilation = 1
if groups is None: groups = 1
layer = [
init_default(
nn.Conv2d(cin,
cout,
ksize,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias), nn.init.kaiming_normal_)
]
if use_bn: layer.append(batchnorm_2d(cout, norm_type=bn))
if use_relu: layer.append(relu(True))
return nn.Sequential(*layer)
def head_layer(self):
cin = self.channels #* 2
return nn.Sequential(
nn.AdaptiveAvgPool2d(1), #AdaptiveConcatPool2d(1),
#nn.Dropout2d(0.5),
Flatten(),
init_default(nn.Linear(cin, self.classes),
nn.init.kaiming_normal_))
def stem_blk(cin,
cout=None,
ksize=3,
stride=1,
use_relu=True,
use_bn=True,
bn=NormType.Batch,
bias=False,
pool='avg'):
if cout is None: cout = cin
padding = ksize // 2
layer = [
init_default(
nn.Conv2d(cin,
cout,
ksize,
stride=stride,
padding=padding,
bias=bias), nn.init.kaiming_normal_)
]
if use_bn: layer.append(batchnorm_2d(cout, norm_type=bn))
if use_relu: layer.append(relu(True))
if pool == 'max': layer.append(nn.MaxPool2d(2, stride=2))
if pool == 'avg':
layer.append(
nn.AvgPool2d(2, stride=2, ceil_mode=True, count_include_pad=False))
layer.append(
init_default(
nn.Conv2d(cout,
cout * 2,
ksize,
stride=stride,
padding=padding,
bias=bias), nn.init.kaiming_normal_))
if use_bn: layer.append(batchnorm_2d(cout * 2, norm_type=bn))
if use_relu: layer.append(relu(True))
if pool == 'max': layer.append(nn.MaxPool2d(2, stride=2))
if pool == 'avg':
layer.append(
nn.AvgPool2d(2, stride=2, ceil_mode=True, count_include_pad=False))
return nn.Sequential(*layer)
def stem(self):
return nn.Sequential(init_default(nn.Conv2d(3, self.channels, 3, padding=1, bias=False),
nn.init.kaiming_normal_),
nn.BatchNorm2d(self.channels))