def __init__(self, p=0.1):
super().__init__()
self.dropout = TN.Dropout2d(p=p)
self.use_dropout = True
if 'torchseg_use_dropout' in os.environ.keys():
self.use_dropout = str2bool(os.environ['torchseg_use_dropout'])
def __init__(self,
in_channels,
out_channels,
kernel_size=1,
padding=0,
stride=1,
eps=1e-5,
momentum=0.1,
inplace=False):
"""
out_channels: class number
upsample_ratio: 2**upsample_layer
"""
super().__init__()
bias = str2bool(os.environ['torchseg_use_bias'])
self.conv_bn_relu = TN.Sequential(
TN.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
padding=padding,
stride=stride,
bias=bias),
local_bn(num_features=out_channels, eps=eps, momentum=momentum),
TN.ReLU(inplace=inplace), local_dropout2d(p=0.1))
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self, in_channels, out_channels, output_shape):
super().__init__()
self.output_shape = output_shape
bias = str2bool(os.environ['torchseg_use_bias'])
self.conv_seq = TN.Sequential(
TN.Conv2d(in_channels=in_channels,
out_channels=4096,
kernel_size=7,
stride=1,
padding=0,
bias=bias), TN.ReLU(), local_dropout2d(p=0.5),
TN.Conv2d(in_channels=4096,
out_channels=4096,
kernel_size=1,
stride=1,
padding=0,
bias=bias), TN.ReLU(), local_dropout2d(p=0.5))
self.conv = TN.Conv2d(in_channels=4096,
out_channels=out_channels,
kernel_size=1,
padding=0,
stride=1,
bias=bias)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def make_layers(cfg,
batch_norm=False,
group_norm=False,
eps=1e-5,
momentum=0.1,
use_none_layer=None,
in_channels=3):
layers = []
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
elif v == 'N':
if use_none_layer is None:
if 'use_none_layer' in os.environ.keys():
use_none_layer = str2bool(os.environ['use_none_layer'])
else:
warnings.warn('use default value for use_none_layer')
use_none_layer = True
if use_none_layer:
layers += [NoneLayer()]
else:
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
if group_norm:
assert not batch_norm, 'group norm will overwrite batch_norm'
conv2d = nn.Conv2d(in_channels,
v,
kernel_size=3,
padding=1,
bias=False)
layers += [
conv2d,
nn.GroupNorm(num_groups=32, num_channels=v, eps=eps),
nn.ReLU(inplace=True)
]
elif batch_norm:
conv2d = nn.Conv2d(in_channels,
v,
kernel_size=3,
padding=1,
bias=False)
layers += [
conv2d,
nn.BatchNorm2d(v, eps=eps, momentum=momentum),
nn.ReLU(inplace=True)
]
else:
conv2d = nn.Conv2d(in_channels,
v,
kernel_size=3,
padding=1,
bias=True)
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
def __init__(self, num_features, eps=1e-5, momentum=0.1):
super().__init__()
self.bn = TN.BatchNorm2d(num_features=num_features,
eps=eps,
momentum=momentum,
affine=True)
self.use_bn = True
if 'torchseg_use_bn' in os.environ.keys():
self.use_bn = str2bool(os.environ['torchseg_use_bn'])
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self, dilation_sizes, class_number, eps=1e-5, momentum=0.1):
"""
in_channels=class_number
out_channels=class_number
backbone->mid_net->global_net->upsample
"""
super(transform_global, self).__init__()
dil_paths = []
bias = str2bool(os.environ['torchseg_use_bias'])
for dilation_size in dilation_sizes:
# for stride=1, to keep size: 2*padding=dilation*(kernel_size-1)
seq = TN.Sequential(
TN.Conv2d(in_channels=class_number,
out_channels=class_number,
kernel_size=3,
stride=dilation_size,
padding=dilation_size,
bias=bias),
local_bn(num_features=class_number, eps=eps,
momentum=momentum), TN.ReLU(), local_dropout2d(p=0.1))
dil_paths.append(seq)
self.dil_paths = TN.ModuleList(dil_paths)
self.conv = TN.Conv2d(in_channels=class_number *
(1 + len(dilation_sizes)),
out_channels=class_number,
kernel_size=1,
padding=0,
bias=bias)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self,
in_channels,
out_channels,
output_shape,
eps=1e-5,
momentum=0.1):
"""
out_channels: class number
"""
super().__init__()
self.output_shape = output_shape
self.center_channels = 512
self.conv_bn_relu = conv_bn_relu(in_channels=in_channels,
out_channels=512,
kernel_size=3,
stride=1,
padding=1,
eps=eps,
momentum=momentum)
# for single stream network, the last conv not need bias?
bias = str2bool(os.environ['torchseg_use_bias'])
self.conv = TN.Conv2d(in_channels=512,
out_channels=out_channels,
kernel_size=1,
padding=0,
stride=1,
bias=bias)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self,
block,
layers,
modify_resnet_head=None,
momentum=0.1,
upsample_layer=5,
in_channels=3,
use_none_layer=True):
self.inplanes = 64
self.momentum = momentum
self.in_channels = in_channels
self.use_none_layer = use_none_layer
super(ResNet, self).__init__()
# for pspnet, the layer1_in_channels=128
# but from checkpoint, the layer1_in_channels=64, make layer1 unchanged!!!
self.layer1_in_channels = 64
self.upsample_layer = upsample_layer
# must set the environment variable before!!!
if modify_resnet_head is None:
warnings.warn('use config from os.environ[modify_resnet_head]')
self.modify_resnet_head = str2bool(
os.environ['modify_resnet_head'])
else:
self.modify_resnet_head = modify_resnet_head
if self.modify_resnet_head:
self.prefix_net = nn.Sequential(
self.conv_bn_relu(in_channels=in_channels,
out_channels=64,
kernel_size=3,
stride=2,
padding=1),
self.conv_bn_relu(in_channels=64,
out_channels=64,
kernel_size=3,
stride=1,
padding=1),
self.conv_bn_relu(in_channels=64,
out_channels=self.layer1_in_channels,
kernel_size=3,
stride=1,
padding=1),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
else:
self.conv1 = nn.Conv2d(in_channels,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64, momentum=momentum)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.prefix_net = nn.Sequential(self.conv1, self.bn1, self.relu,
self.maxpool)
self.layer1 = self._make_layer(block,
64,
layers[0],
index=1,
momentum=momentum)
self.layer2 = self._make_layer(block,
128,
layers[1],
index=2,
stride=2,
momentum=momentum)
self.layer3 = self._make_layer(block,
256,
layers[2],
index=3,
stride=2,
momentum=momentum)
self.layer4 = self._make_layer(block,
512,
layers[3],
index=4,
stride=2,
momentum=momentum)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def __init__(self,
output_stride,
input_shape,
out_channels,
eps=1e-05,
momentum=0.1):
"""
output_stride=2**upsample_layer
input_shape=b,c,h,w
output_shape=b,out_channels,h,w
"""
super().__init__()
b, in_channels, h, w = input_shape
assert h == w, 'height=%d and width=%d must equal in input_shape' % (h,
w)
if output_stride not in [8, 16]:
raise ValueError('output_stride must be either 8 or 16.')
atrous_rates = [1, 6, 12, 18]
if output_stride == 8:
atrous_rates = [2 * rate for rate in atrous_rates]
atrous_rates[0] = 1
out_c = out_channels
self.out_channels = out_channels
atrous_paths = []
k_sizes = [1, 3, 3, 3]
bias = str2bool(os.environ['torchseg_use_bias'])
# in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True
for r, k in zip(atrous_rates, k_sizes):
path = TN.Sequential(
TN.Conv2d(in_channels=in_channels,
out_channels=out_c,
kernel_size=k,
stride=1,
padding=r * (k - 1) // 2,
dilation=r,
bias=bias),
local_bn(num_features=out_c, eps=eps, momentum=momentum),
TN.ReLU())
atrous_paths.append(path)
self.image_level_path = TN.Sequential(
TN.AvgPool2d(kernel_size=h),
TN.Conv2d(in_channels=in_channels,
out_channels=out_c,
kernel_size=1,
stride=1,
padding=0,
bias=bias),
local_bn(num_features=out_c, eps=eps, momentum=momentum),
TN.ReLU(),
local_upsample(size=(h, w), mode='bilinear', align_corners=True))
self.final_conv = TN.Sequential(
TN.Conv2d(in_channels=out_c * 5,
out_channels=out_c,
kernel_size=1,
stride=1,
padding=0,
bias=bias),
local_bn(num_features=out_c, eps=eps, momentum=momentum),
TN.ReLU())
self.atrous_paths = TN.ModuleList(atrous_paths)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self,
in_channels,
depth,
class_number,
fusion_type='max',
do_fusion=False,
eps=1e-5,
momentum=0.1):
"""
input [b,in_channels,h,w]
output [b,class_number,h,w]
"""
super().__init__()
self.depth = depth
self.class_number = class_number
self.fusion_type = fusion_type
self.do_fusion = do_fusion
if do_fusion:
assert depth > 1, 'fusion can only do once on the top level'
support_types = ['max', 'mean', 'route']
assert fusion_type in support_types, 'fusion_type %s not in %s' % (
fusion_type, support_types)
assert depth >= 1, 'depth %d cannot less than 1' % depth
fractal_paths = []
if fusion_type == 'route':
raise NotImplementedError
bias = str2bool(os.environ['torchseg_use_bias'])
if depth == 1:
path = TN.Sequential(
TN.Conv2d(in_channels=in_channels,
out_channels=class_number,
kernel_size=1,
stride=1,
padding=0,
bias=bias),
local_bn(num_features=class_number, eps=eps,
momentum=momentum), TN.ReLU())
fractal_paths.append(path)
else:
for i in range(depth):
if i == 0:
path = transform_fractal(in_channels, i + 1, class_number,
fusion_type)
else:
path = TN.Sequential(
TN.Conv2d(in_channels=in_channels,
out_channels=(2**i) * class_number,
kernel_size=1,
stride=1,
padding=0,
bias=bias),
local_bn(num_features=(2**i) * class_number),
TN.ReLU(),
transform_fractal((2**i) * class_number, i,
class_number, fusion_type))
fractal_paths.append(path)
self.fractal_paths = TN.ModuleList(fractal_paths)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)
def __init__(self,
pool_sizes,
scale,
input_shape,
out_channels,
eps=1e-5,
momentum=0.1):
"""
pool_sizes = [1,2,3,6]
scale = 5,10
out_channels = the output channel for transform_psp
out_size = ?
"""
super(transform_psp, self).__init__()
self.input_shape = input_shape
b, in_channels, height, width = input_shape
path_out_c_list = []
N = len(pool_sizes)
assert out_channels > in_channels, 'out_channels will concat inputh, so out_chanels=%d should >= in_chanels=%d' % (
out_channels, in_channels)
# the output channel for pool_paths
pool_out_channels = out_channels - in_channels
mean_c = pool_out_channels // N
for i in range(N - 1):
path_out_c_list.append(mean_c)
path_out_c_list.append(pool_out_channels + mean_c - mean_c * N)
self.pool_sizes = pool_sizes
self.scale = scale
pool_paths = []
bias = str2bool(os.environ['torchseg_use_bias'])
for pool_size, out_c in zip(pool_sizes, path_out_c_list):
pool_path = TN.Sequential(
TN.AvgPool2d(kernel_size=pool_size * scale,
stride=pool_size * scale,
padding=0),
TN.Conv2d(in_channels=in_channels,
out_channels=out_c,
kernel_size=1,
stride=1,
padding=0,
bias=bias),
local_bn(num_features=out_c, eps=eps, momentum=momentum),
TN.ReLU(), local_dropout2d(p=0.1),
local_upsample(size=(height, width),
mode='bilinear',
align_corners=True))
pool_paths.append(pool_path)
self.pool_paths = TN.ModuleList(pool_paths)
for m in self.modules():
if isinstance(m, TN.Conv2d):
TN.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, TN.BatchNorm2d):
TN.init.constant_(m.weight, 1)
TN.init.constant_(m.bias, 0)