def __init__(self,
encoder: nn.Module,
n_classes,
final_bias=0.,
chs=256,
n_anchors=9,
flatten=True):
super().__init__()
self.n_classes, self.flatten = n_classes, flatten
imsize = (256, 256)
sfs_szs = model_sizes(encoder, size=imsize)
sfs_idxs = list(reversed(_get_sz_change_idxs(sfs_szs)))
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
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, sfs_szs[idx][1], hook)
for idx, hook in zip(sfs_idxs[-2:-4:-1], self.sfs[-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)
def __init__(self,
encoder,
n_classes,
img_size,
blur=False,
blur_final=True,
self_attention=False,
y_range=None,
bottle=False,
act_cls=defaults.activation,
init=nn.init.kaiming_normal_,
norm_type=None,
include_encoder=True,
include_middle_conv=True,
**kwargs):
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()
layers = []
if include_encoder:
layers.append(encoder)
if include_middle_conv:
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 += [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 = int(x.shape[1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i == len(sz_chg_idxs) - 3)
noskip_unet_block = NoSkipUnetBlock(up_in_c,
final_div=not_final,
blur=do_blur,
self_attention=sa,
act_cls=act_cls,
init=init,
norm_type=norm_type,
**kwargs).eval()
layers.append(noskip_unet_block)
x = noskip_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 += [
ConvLayer(ni,
n_classes,
ks=1,
act_cls=None,
norm_type=norm_type,
**kwargs)
]
if include_middle_conv:
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)