def multi_task_unet_learner(*args, log_vars=None, **kwargs):
"""
Creates a learner suited for classification+segmentation multii-task
learning problem
args: positional arguments for cnn_learner and unet_learner
kwargs: keayword arguments for cnn_learner and unet_learner
return: learner that contains MultiTaskModel
"""
unet_learn = unet_learner(*args, **kwargs)
sfs_idxs = unet_learn.model.sfs_idxs
cnn_learn = cnn_learner(*args, **kwargs)
base = unet_learn.model[0]
unet_head = unet_learn.model[1:]
hooks = hook_outputs([base[i] for i in sfs_idxs])
for block, hook in zip(unet_head[3:7], hooks):
block.hook = hook
heads = [cnn_learn.model[1:], unet_head]
unet_learn.model = MultiTaskModel(base, heads, log_vars=log_vars).to(
unet_learn.data.device)
lg = unet_learn.layer_groups
lg[2] = nn.Sequential(*list(lg[2]), *flatten_model(heads[0]),
unet_learn.model.log_vars)
unet_learn.layer_groups = lg
unet_learn.create_opt(slice(1e-3))
return unet_learn
def __init__(self,
encoder: nn.Module,
n_classes: int,
blur: bool = False,
blur_final=True,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
**kwargs):
imsize = (args.size, args.size)
sfs_szs = model_sizes(encoder, size=imsize)
sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
x = dummy_eval(encoder, imsize).detach()
ni = sfs_szs[-1][1]
middle_conv = nn.Sequential(conv_layer(ni, ni * 2, **kwargs),
conv_layer(ni * 2, ni, **kwargs)).eval()
x = middle_conv(x)
layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
self.hc_hooks = [Hook(layers[-1], _hook_inner, detach=False)]
hc_c = [x.shape[1]]
for i, idx in enumerate(sfs_idxs):
not_final = i != len(sfs_idxs) - 1
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i == len(sfs_idxs) - 3)
unet_block = UnetBlock(up_in_c,
x_in_c,
self.sfs[i],
final_div=not_final,
blur=blur,
self_attention=sa,
**kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
self.hc_hooks.append(Hook(layers[-1], _hook_inner, detach=False))
hc_c.append(x.shape[1])
ni = x.shape[1]
if imsize != sfs_szs[0][-2:]:
layers.append(PixelShuffle_ICNR(ni, **kwargs))
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(res_block(ni, bottle=bottle, **kwargs))
hc_c.append(ni)
layers.append(Hcolumns(self.hc_hooks, hc_c))
layers += [
conv_layer(ni * len(hc_c),
n_classes,
ks=1,
use_activ=False,
**kwargs)
]
if y_range is not None: layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
def __init__(self, encoder, n_classes, final_bias=0., chs=256, n_anchors=9, flatten=True, chip_size=(256,256), n_bands=3):
# chs - channels for top down layers in FPN
super().__init__()
self.n_classes,self.flatten = n_classes,flatten
self.chip_size = chip_size
# Fetch the sizes of various activation layers of the backbone
sfs_szs = model_sizes(encoder, size=self.chip_size)
hooks = hook_outputs(encoder)
self.encoder = encoder
self.c5top5 = conv2d(sfs_szs[-1][1], chs, ks=1, bias=True)
self.c5top6 = conv2d(sfs_szs[-1][1], chs, stride=2, bias=True)
self.p6top7 = nn.Sequential(nn.ReLU(), conv2d(chs, chs, stride=2, bias=True))
self.merges = nn.ModuleList([LateralUpsampleMerge(chs, szs[1], hook)
for szs,hook in zip(sfs_szs[-2:-4:-1], hooks[-2:-4:-1])])
self.smoothers = nn.ModuleList([conv2d(chs, chs, 3, bias=True) for _ in range(3)])
self.classifier = self._head_subnet(n_classes, n_anchors, final_bias, chs=chs)
self.box_regressor = self._head_subnet(4, n_anchors, 0., chs=chs)
# Create a dummy x to be passed through the model and fetch the sizes
x_dummy = torch.rand(n_bands,self.chip_size[0],self.chip_size[1]).unsqueeze(0)
p_states = self._create_p_states(x_dummy)
self.sizes = [[p.size(2), p.size(3)] for p in p_states]
def __init__(self, m_feat, layer_ids, layer_wgts):
super().__init__()
self.m_feat = m_feat
self.base_loss = F.l1_loss
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = ([
"pixel",
] + [f"feat_{i}" for i in range(len(layer_ids))] +
[f"gram_{i}" for i in range(len(layer_ids))])
def __init__(self, m_feat, layer_ids, layer_wgts,mrse=True):
super().__init__()
self.m_feat = m_feat
self.loss_features = [self.m_feat[i] for i in layer_ids]
self.hooks = hook_outputs(self.loss_features, detach=False)
self.wgts = layer_wgts
self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))
] + [f'gram_{i}' for i in range(len(layer_ids))]
self.mrae = MRAELoss()
self.mrse = MRSELoss()
self.mse = MSELossFlat()
self.mrse_switch = mrse
def __init__(self,
encoder=None,
n_classes=2,
last_filters=32,
imsize=(256, 256),
y_range=None,
**kwargs):
self.n_classes = n_classes
layers = nn.ModuleList()
# Encoder
sfs_szs = model_sizes(encoder, size=imsize)
sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
layers.append(encoder)
x = dummy_eval(encoder, imsize).detach()
self.hc_hooks = []
hc_c = []
ni = sfs_szs[-1][1]
middle_conv = nn.Sequential(conv_layer(ni, ni * 2),
conv_layer(ni * 2, ni)).eval()
x = middle_conv(x)
layers.extend([batchnorm_2d(ni), nn.ReLU(), middle_conv])
# self.hc_hooks = [Hook(layers[-1], _hook_inner, detach=False)]
# hc_c = [x.shape[1]]
# Decoder
n_filters = [64, 128, 256, 512]
n = len(n_filters)
is_deconv = True
for i, idx in enumerate(sfs_idxs[:-1]):
in_c, out_c = int(n_filters[n - i - 1] +
n_filters[n - i - 2]) // 2, int(sfs_szs[idx][1])
dec_bloc = DecoderBlock(in_c, out_c, self.sfs[i], is_deconv,
True).eval()
layers.append(dec_bloc)
x = dec_bloc(x)
self.hc_hooks.append(Hook(layers[-1], _hook_inner, detach=False))
hc_c.append(x.shape[1])
ni = x.shape[1]
layers.append(PixelShuffle_ICNR(n_filters[0], scale=2))
layers.append(Hcolumns(self.hc_hooks, hc_c))
fin_block = FinalBlock(ni * (len(hc_c) + 1), last_filters, n_classes)
layers.append(fin_block)
if y_range is not None:
layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
def __init__(self,
encoder: nn.Module,
n_classes: int,
blur: bool = False,
blur_final=True,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
norm_type: Optional[NormType] = NormType.Batch,
nf_factor: int = 1,
**kwargs):
nf = 512 * nf_factor
extra_bn = norm_type == NormType.Spectral
imsize = (256, 256)
sfs_szs = model_sizes(encoder, size=imsize)
sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs], detach=False)
x = dummy_eval(encoder, imsize).detach()
ni = sfs_szs[-1][1]
middle_conv = nn.Sequential(
custom_conv_layer(ni,
ni * 2,
norm_type=norm_type,
extra_bn=extra_bn,
**kwargs),
custom_conv_layer(ni * 2,
ni,
norm_type=norm_type,
extra_bn=extra_bn,
**kwargs),
).eval()
x = middle_conv(x)
layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
for i, idx in enumerate(sfs_idxs):
not_final = i != len(sfs_idxs) - 1
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i == len(sfs_idxs) - 3)
n_out = nf if not_final else nf // 2
unet_block = UnetBlockWide(up_in_c,
x_in_c,
n_out,
self.sfs[i],
final_div=not_final,
blur=blur,
self_attention=sa,
norm_type=norm_type,
extra_bn=extra_bn,
**kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
ni = x.shape[1]
if imsize != sfs_szs[0][-2:]:
layers.append(PixelShuffle_ICNR(ni, **kwargs))
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(
res_block(ni, bottle=bottle, norm_type=norm_type, **kwargs))
layers += [
custom_conv_layer(ni,
n_classes,
ks=1,
use_activ=False,
norm_type=norm_type)
]
if y_range is not None:
layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
n = len(self.hooks)
out = [F.interpolate(self.hooks[i].stored if self.factorization is None
else self.factorization[i](self.hooks[i].stored), scale_factor=2**(self.n-i),
mode='bilinear',align_corners=False) for i in range(self.n)] + [x]
return torch.cat(out, dim=1)
class DynamicUnet_Hcolumns(SequentialEx):
"Create a U-Net from a given architecture."
def __init__(self, encoder:nn.Module, n_classes:int, blur:bool=False, blur_final=True,
self_attention:bool=False,
y_range:Optional[Tuple[float,float]]=None,
last_cross:bool=True, bottle:bool=False, **kwargs):
imsize = (args.size, args.size)
sfs_szs = model_sizes(encoder, size=imsize)
sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
x = dummy_eval(encoder, imsize).detach()
ni = sfs_szs[-1][1]
middle_conv = nn.Sequential(conv_layer(ni, ni*2, **kwargs),
conv_layer(ni*2, ni, **kwargs)).eval()
x = middle_conv(x)
layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
self.hc_hooks = [Hook(layers[-1], _hook_inner, detach=False)]
hc_c = [x.shape[1]]
for i,idx in enumerate(sfs_idxs):
not_final = i!=len(sfs_idxs)-1
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
do_blur = blur and (not_final or blur_final)
def __init__(
self,
encoder: nn.Module,
n_classes: int,
blur: bool = False,
blur_final=True,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
small=True,
**kwargs,
):
imsize = (256, 256)
# for resnet50 ... but memory not enough...
# sfs_szs = [(1, 64, 128, 128), (1, 64, 128, 128), (1, 64, 1...512, 32, 32), (1, 1024, 16, 16), (1, 2048, 8, 8)]
# sfs_idxs = [6, 5, 4, 2] #? 3?
sfs_szs = model_sizes(encoder, size=imsize)
# for resnext50_32x4d
# [torch.Size([1, 64, 64, 64]), torch.Size([1, 64, 64, 64]), torch.Size([1, 64, 64, 64]), torch.Size([1, 64, 32, 32]), torch.Size([1, 256, 32, 32]), torch.Size([1, 512, 16, 16]), torch.Size([1, 1024, 8, 8]), torch.Size([1, 2048, 4, 4])]
sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
# if small: sfs_idxs = sfs_idxs[-3:] (need to do double upscale)
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
x = dummy_eval(encoder, imsize).detach()
ni = sfs_szs[-1][1]
if small:
middle_conv_size_down_scale = 2
middle_conv = conv_layer(ni, ni // middle_conv_size_down_scale,
**kwargs).eval()
else:
middle_conv_size_scale = 2
middle_conv = nn.Sequential(
conv_layer(ni, ni * middle_conv_size_scale, **kwargs),
conv_layer(ni * middle_conv_size_scale, ni, **kwargs),
).eval()
x = middle_conv(x)
layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
if small:
self.hc_hooks = []
hc_c = []
else:
self.hc_hooks = [Hook(layers[-1], _hook_inner, detach=False)]
hc_c = [x.shape[1]]
for i, idx in enumerate(sfs_idxs):
final_unet_flag = i == len(sfs_idxs) - 1
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
do_blur = blur and (final_unet_flag or blur_final)
sa = self_attention and (i == len(sfs_idxs) - 3)
unet_block_class = UnetBlockSmall if small else UnetBlock
unet_block = unet_block_class(
up_in_c,
x_in_c,
self.sfs[i],
final_div=final_unet_flag,
blur=blur,
self_attention=sa,
**kwargs,
).eval()
print(unet_block)
layers.append(unet_block)
x = unet_block(x)
# added for hypercolumns, two line
self.hc_hooks.append(Hook(layers[-1], _hook_inner, detach=False))
hc_c.append(x.shape[1])
ni = x.shape[1]
if imsize != sfs_szs[0][-2:]:
layers.append(PixelShuffle_ICNR(ni, **kwargs))
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(res_block(ni, bottle=bottle, **kwargs))
# added for hypercolumns, two line
hc_c.append(ni)
layers.append(Hcolumns(self.hc_hooks, hc_c))
layers += [
conv_layer(ni * len(hc_c),
n_classes,
ks=1,
use_activ=False,
**kwargs)
]
if y_range is not None:
layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
def __init__(self, encoder=None, n_classes=2, last_filters=32, imsize=(256, 256), y_range=None, **kwargs):
self.n_classes = n_classes
layers = nn.ModuleList()
# Encoder
sfs_idxs = [4, 3, 2, 1, 0]
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
layers.append(encoder)
x = dummy_eval(encoder, imsize).detach()
self.hc_hooks = []
hc_c = []
# ni = sfs_szs[-1][1]
# middle_conv = nn.Sequential(conv_layer(ni, ni * 2),
# conv_layer(ni * 2, ni)).eval()
# x = middle_conv(x)
# layers.extend([batchnorm_2d(ni), nn.ReLU(), middle_conv])
# Decoder
n_filters = [128, 256, 512, 1024, 2048]
n = len(n_filters)
is_deconv = True
for i, idx in enumerate(sfs_idxs[:-1]):
if i == 0:
in_c, out_c = n_filters[n - i - 1], n_filters[n - i - 2]
else:
in_c, out_c = 2 * n_filters[n - i - 1], n_filters[n - i - 2]
if i == 3:
scale = False
else:
scale = True
dec_bloc = DecoderBlock(in_c, out_c, self.sfs[i+1], is_deconv, scale).eval()
layers.append(dec_bloc)
x = dec_bloc(x)
self.hc_hooks.append(Hook(layers[-1], _hook_inner, detach=False))
hc_c.append(x.shape[1])
# print(x.size())
# last decoder
in_c, out_c = n_filters[0] + 64, n_filters[0]
dec_bloc = DecoderBlock(in_c, out_c, None, is_deconv, True).eval()
layers.append(dec_bloc)
x = dec_bloc(x)
self.hc_hooks.append(Hook(layers[-1], _hook_inner, detach=False))
hc_c.append(x.shape[1])
# print(x.size())
ni = x.shape[1]
layers.append(PixelShuffle_ICNR(n_filters[0], scale=2))
layers.append(Hcolumns(self.hc_hooks, hc_c))
fin_block = FinalBlock(ni * (len(hc_c) + 1), last_filters, n_classes)
layers.append(fin_block)
if y_range is not None:
layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
