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,
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_sfs_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: nn.Module,
n_classes,
final_bias: float = 0.,
n_conv: float = 4,
chs=256,
n_anchors=9,
anchors=None,
oversize=1.5,
flatten=True,
sizes=None,
patchClassifier: nn.Module = None):
super().__init__()
self.n_classes, self.flatten = n_classes, flatten
imsize = (256, 256)
self.sizes = sizes
self.grid_target_size = (64, 64)
self._sizes = torch.Tensor(sizes)
self._oversize = oversize
sfs_szs, x, hooks = self._model_sizes(encoder, size=imsize)
sfs_idxs = _get_sfs_idxs(sfs_szs)
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,
n_conv=n_conv)
self.box_regressor = self._head_subnet(3,
n_anchors,
0.,
chs=chs,
n_conv=n_conv)
if (patchClassifier is None):
self.box_stn = create_cnn_model(models.resnet18,
n_classes,
cut=None,
pretrained=True,
lin_ftrs=None,
ps=0.5,
split_on=None,
bn_final=False,
concat_pool=True)
# self._stn_subnet(n_classes, final_bias, chs=chs, n_conv=n_conv)
else:
self.box_stn = patchClassifier
self.anchors = anchors
def __init__(self,
vol_size,
encoder: nn.Module,
n_classes: int,
last_cross: bool = True,
bottle: bool = False,
pixel_shuffle=False,
light_up_block=True,
self_attention: bool = False,
**kwargs):
sfs_szs = model_sizes(encoder, size=vol_size)
# sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
sfs_idxs = np.unique(_get_sfs_idxs(sfs_szs))[::-1].tolist()
self.sfs = hook_outputs([encoder[i] for i in sfs_idxs])
x = dummy_eval(encoder, vol_size).detach()
ni = sfs_szs[-1][1]
middle_conv = nn.Sequential(conv_layer3d(ni, ni * 2, **kwargs),
conv_layer3d(ni * 2, ni, **kwargs)).eval()
x = middle_conv(x)
layers = [encoder, batchnorm_3d(ni), nn.ReLU(), middle_conv]
for i, idx in enumerate(sfs_idxs):
not_final = i != len(sfs_idxs) - 1
ds = not any(
[a >= b for a, b in zip(x.shape[-3:], sfs_szs[idx][-3:])])
ds = ds if pixel_shuffle else pixel_shuffle
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
sa = self_attention and (i == len(sfs_idxs) - 3)
vnet_block = VnetBlock(up_in_c,
x_in_c,
self.sfs[i],
final_div=not_final,
do_shuffle=ds,
double_out_conv=not light_up_block,
self_attention=sa,
**kwargs).eval()
layers.append(vnet_block)
x = vnet_block(x)
ni = x.shape[1]
if vol_size != sfs_szs[0][-3:]:
#layers.append(PixelShuffle_ICNR(ni, **kwargs))
layers.append(Upsample3d(vol_size))
if last_cross:
layers.append(MergeLayer(dense=True))
ni += in_channels(encoder)
layers.append(res_block3d(ni, bottle=bottle, **kwargs))
layers += [
conv_layer3d(ni, n_classes, ks=1, use_activ=False, **kwargs)
]
super().__init__(*layers)
def __init__(self, encoder: nn.Module, n_classes, final_bias:float=0., n_conv:float=4,
chs=512, n_anchors=9, flatten=True, sizes=None):
super().__init__()
self.n_classes, self.flatten = n_classes, flatten
imsize = (512, 512)
self.sizes = sizes
sfs_szs, x, hooks = self._model_sizes(encoder, size=imsize)
sfs_idxs = _get_sfs_idxs(sfs_szs)
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, n_conv=n_conv)
self.box_regressor = self._head_subnet(4, n_anchors, 0., chs=chs, n_conv=n_conv)
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,
y_range:Optional[Tuple[float,float]]=None, skip_connections=True,
**kwargs):
encoder[2] = encoder[0:3]
encoder = nn.Sequential(*list(encoder.children())[2:])
attented_layers = []
filter = []
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])
sfs_idxs = sfs_idxs[:-1]
attented_layers.extend([encoder[ind] for ind in sfs_idxs[::-1]])
attented_layers.append(encoder[-1])
filter.extend([sfs_szs[ind][1] for ind in sfs_idxs[::-1]])
x = dummy_eval(encoder, imsize).detach()
ni = sfs_szs[-1][1]
filter.append(ni)
middle_conv_enc = conv_layer(ni, ni*2, **kwargs).eval()
middle_conv_dec = conv_layer(ni*2, ni, **kwargs).eval()
x = middle_conv_enc(x)
x = middle_conv_dec(x)
layers = list(encoder)
layers = layers + [batchnorm_2d(ni), nn.ReLU(), middle_conv_enc, middle_conv_dec]
attented_layers.append(middle_conv_enc)
attented_layers.append(middle_conv_dec)
filter.extend([ni*2,ni*2,ni])
# sfs_idxs = sfs_idxs[:-2]
for i,idx in enumerate(sfs_idxs):
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
if skip_connections:
not_final = not (i!=len(sfs_idxs)-1)
unet_block = UnetBlock(up_in_c, x_in_c, self.sfs[i], final_div=not_final,
**kwargs).eval()
else:
unet_block = UnetBlockWithoutSkipConnection(up_in_c, **kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
attented_layers.append(layers[-1])
filter.append(x_in_c) # in for first filter param for attention block
filter = filter[:-1]
ni = x.shape[1]
unet_block_last = UnetBlockWithoutSkipConnection(10,
# final_div=not_final,
blur=False, self_attention=False,
**kwargs)
if imsize != sfs_szs[0][-2:]:
unet_block_last.shuf = PixelShuffle_ICNR(ni, **kwargs)
else:
unet_block_last.shuf = nn.Identity()
unet_block_last.conv1 = conv_layer(ni, n_classes, ks=1, use_activ=False, **kwargs)
unet_block_last.conv2 = nn.Identity()
unet_block_last.relu = nn.Identity()
layers.append(unet_block_last)
attented_layers.append(unet_block_last)
# if skip_connections:
# ni = 32
filter.extend([ni, n_classes])
if y_range is not None: layers.append(SigmoidRange(*y_range))
super().__init__(*layers)
self.attended_layers = attented_layers
self.filter = filter
def forward(self, x:Tensor):
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])
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)
