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: nn.Module, n_classes: int, img_size: Tuple[int, int] = (256, 256),
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 = tuple(img_size)
sfs_szs, select_layer = get_unet_config(encoder, img_size)
ni = sfs_szs[-1][1]
sfs_szs = list(reversed(sfs_szs[:-1]))
select_layer = list(reversed(select_layer[:-1]))
self.sfs = hook_outputs(select_layer, detach=False)
x = dummy_eval(encoder, imsize).detach()
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]
for i, x_size in enumerate(sfs_szs):
not_final = i != len(sfs_szs) - 1
up_in_c, x_in_c = int(x.shape[1]), int(x_size[1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i == len(sfs_szs) - 3)
unet_block = UnetBlock(up_in_c, x_in_c, self.sfs[i], final_div=not_final, blur=do_blur, self_attention=sa,
**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))
x = PixelShuffle_ICNR(ni)(x)
if imsize != x.shape[-2:]: layers.append(Lambda(lambda x: F.interpolate(x, imsize, mode='nearest')))
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(res_block(ni, bottle=bottle, **kwargs))
layers += [conv_layer(ni, n_classes, ks=1, use_activ=False, **kwargs)]
if y_range is not None: layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
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))
def __init__(self,
arch=resnet50,
n_classes=32,
img_size=(96, 128),
blur=False,
blur_final=True,
y_range=None,
last_cross=True,
bottle=False,
init=nn.init.kaiming_normal_,
norm_type=None,
self_attention=None,
act_cls=defaults.activation,
n_in=3,
cut=None,
**kwargs):
meta = model_meta.get(arch, _default_meta)
encoder = create_body(arch,
n_in,
pretrained=False,
cut=ifnone(cut, meta["cut"]))
imsize = img_size
sizes = model_sizes(encoder, size=imsize)
sz_chg_idxs = list(reversed(_get_sz_change_idxs(sizes)))
self.sfs = hook_outputs([encoder[i] for i in sz_chg_idxs],
detach=False)
x = dummy_eval(encoder, imsize).detach()
ni = sizes[-1][1]
middle_conv = nn.Sequential(
ConvLayer(ni,
ni * 2,
act_cls=act_cls,
norm_type=norm_type,
**kwargs),
ConvLayer(ni * 2,
ni,
act_cls=act_cls,
norm_type=norm_type,
**kwargs),
).eval()
x = middle_conv(x)
layers = [encoder, BatchNorm(ni), nn.ReLU(), middle_conv]
for i, idx in enumerate(sz_chg_idxs):
not_final = i != len(sz_chg_idxs) - 1
up_in_c, x_in_c = int(x.shape[1]), int(sizes[idx][1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i == len(sz_chg_idxs) - 3)
unet_block = UnetBlock(up_in_c,
x_in_c,
self.sfs[i],
final_div=not_final,
blur=do_blur,
self_attention=sa,
act_cls=act_cls,
init=init,
norm_type=norm_type,
**kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
ni = x.shape[1]
if imsize != sizes[0][-2:]:
layers.append(
PixelShuffle_ICNR(ni, act_cls=act_cls, norm_type=norm_type))
layers.append(ResizeToOrig())
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(
ResBlock(1,
ni,
ni // 2 if bottle else ni,
act_cls=act_cls,
norm_type=norm_type,
**kwargs))
layers += [
ConvLayer(ni,
n_classes,
ks=1,
act_cls=None,
norm_type=norm_type,
**kwargs)
]
apply_init(nn.Sequential(layers[3], layers[-2]), init)
# apply_init(nn.Sequential(layers[2]), init)
if y_range is not None:
layers.append(SigmoidRange(*y_range))
super().__init__(*layers)