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, num_classes, backbone_fn, chip_size=224): super().__init__() if getattr(backbone_fn, '_is_multispectral', False): self.backbone = create_body(backbone_fn, pretrained=True, cut=_get_backbone_meta( backbone_fn.__name__)['cut']) else: self.backbone = create_body(backbone_fn, pretrained=True) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(self.backbone.children())[0] else: hookable_modules = list(self.backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' ## Hook at the index where we need to get the auxillary logits out self.hook = hook_output(hookable_modules[modify_dilation_index]) custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), ( padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = ( dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size)) ## Geting the number of channel persent in stored activation inside of the hook num_channels_aux_classifier = self.hook.stored.shape[1] ## Get number of channels in the last layer num_channels_classifier = feature_sizes[-1][1] self.classifier = DeepLabHead(num_channels_classifier, num_classes) self.aux_classifier = FCNHead(num_channels_aux_classifier, num_classes)
def __init__(self, model, chip_size, num_classes): super(AuxPSUnet, self).__init__() self.model = model for idx, i in enumerate(flatten_model(self.model)): if hasattr(i, 'dilation'): dilation = i.dilation dilation = dilation[0] if isinstance(dilation, tuple) else dilation if dilation > 1: break self.hook = hook_output(flatten_model(model)[idx - 1]) ## returns the size of various activations model_sizes(self.model, size=(chip_size, chip_size)) ## Geting the stored parameters inside of the hook aux_in_channels = self.hook.stored.shape[1] del self.hook.stored self.aux_logits = nn.Conv2d(aux_in_channels, num_classes, kernel_size=1)
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)
def cnn_activations_count(model): _, ch, h, w = model_sizes(create_body(models.resnet18), (SZ, SZ))[-1] return ch * h * w
def __init__(self, data, grids=[4, 2, 1], zooms=[0.7, 1., 1.3], ratios=[[1., 1.], [1., 0.5], [0.5, 1.]], backbone=None, drop=0.3, bias=-4., focal_loss=False, pretrained_path=None, location_loss_factor=None): super().__init__() self._device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') # assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)), if location_loss_factor is not None: if not ((location_loss_factor > 0) and (location_loss_factor < 1)): raise Exception( '`location_loss_factor` should be greater than 0 and less than 1' ) self.location_loss_factor = location_loss_factor if not HAS_FASTAI: _raise_fastai_import_error() if backbone is None: self._backbone = models.resnet34 elif type(backbone) is str: self._backbone = getattr(models, backbone) else: self._backbone = backbone self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] ssd_head = SSDHead(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) self._data = data self.learn = create_cnn(data=data, arch=self._backbone, custom_head=ssd_head) self.learn.model = self.learn.model.to(self._device) if pretrained_path is not None: self.load(pretrained_path) if focal_loss: self._loss_f = FocalLoss(data.c) else: self._loss_f = BCE_Loss(data.c) self.learn.loss_func = self._ssd_loss
def _pspnet_unet(num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True): """ Function which returns PPM module attached to backbone which is then used to form the Unet. """ if getattr(backbone_fn, '_is_multispectral', False): backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut']) else: backbone = create_body(backbone_fn, pretrained=pretrained) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(backbone.children())[0] else: hookable_modules = list(backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) # padding = 1 for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(backbone, size=(chip_size, chip_size)) ## Get number of channels in the last layer num_channels = feature_sizes[-1][1] penultimate_channels = num_channels / len(pyramid_sizes) ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes) in_final = int(penultimate_channels) * len(pyramid_sizes) + num_channels # Reduce channel size after pyramid pooling module to avoid CUDA OOM error. final_conv = nn.Conv2d(in_channels=in_final, out_channels=512, kernel_size=3, padding=1) ## To make Dynamic Unet work as it expects a backbone which can be indexed. if "densenet" in backbone_name or "vgg" in backbone_name: backbone = backbone[0] layers = [*backbone, ppm, final_conv] return nn.Sequential(*layers)
def __init__(self, num_classes, backbone_fn, chip_size=224, pyramid_sizes=(1, 2, 3, 6), pretrained=True): super(PSPNet, self).__init__() if getattr(backbone_fn, '_is_multispectral', False): self.backbone = create_body(backbone_fn, pretrained=pretrained, cut=_get_backbone_meta(backbone_fn.__name__)['cut']) else: self.backbone = create_body(backbone_fn, pretrained=pretrained) backbone_name = backbone_fn.__name__ ## Support for different backbones if "densenet" in backbone_name or "vgg" in backbone_name: hookable_modules = list(self.backbone.children())[0] else: hookable_modules = list(self.backbone.children()) if "vgg" in backbone_name: modify_dilation_index = -5 else: modify_dilation_index = -2 if backbone_name == 'resnet18' or backbone_name == 'resnet34': module_to_check = 'conv' else: module_to_check = 'conv2' ## Hook at the index where we need to get the auxillary logits out self.hook = hook_output(hookable_modules[modify_dilation_index]) custom_idx = 0 for i, module in enumerate(hookable_modules[modify_dilation_index:]): dilation = 2 * (i + 1) padding = 2 * (i + 1) for n, m in module.named_modules(): if module_to_check in n: m.dilation, m.padding, m.stride = (dilation, dilation), (padding, padding), (1, 1) elif 'downsample.0' in n: m.stride = (1, 1) if "vgg" in backbone_fn.__name__: if isinstance(module, nn.Conv2d): dilation = 2 * (custom_idx + 1) padding = 2 * (custom_idx + 1) module.dilation, module.padding, module.stride = (dilation, dilation), (padding, padding), (1, 1) custom_idx += 1 ## returns the size of various activations feature_sizes = model_sizes(self.backbone, size=(chip_size, chip_size)) ## Geting the stored parameters inside of the hook aux_in_channels = self.hook.stored.shape[1] ## Get number of channels in the last layer num_channels = feature_sizes[-1][1] penultimate_channels = num_channels / len(pyramid_sizes) self.ppm = _PyramidPoolingModule(num_channels, int(penultimate_channels), pyramid_sizes) self.final = nn.Sequential( ## To handle case when the length of pyramid_sizes is odd nn.Conv2d(int(penultimate_channels) * len(pyramid_sizes) + num_channels, math.ceil(penultimate_channels), kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(math.ceil(penultimate_channels)), nn.ReLU(inplace=True), nn.Dropout(0.1), nn.Conv2d(math.ceil(penultimate_channels), num_classes, kernel_size=1) ) self.aux_logits = nn.Conv2d(aux_in_channels, num_classes, kernel_size=1) initialize_weights(self.aux_logits) initialize_weights(self.ppm, self.final)
def __init__(self, data, grids=None, zooms=[1.], ratios=[[1., 1.]], backbone=None, drop=0.3, bias=-4., focal_loss=False, pretrained_path=None, location_loss_factor=None, ssd_version=2, backend='pytorch'): super().__init__(data, backbone) self._backend = backend if self._backend == 'tensorflow': self._intialize_tensorflow(data, grids, zooms, ratios, backbone, drop, bias, pretrained_path, location_loss_factor) else: # assert (location_loss_factor is not None) or ((location_loss_factor > 0) and (location_loss_factor < 1)), if not ssd_version in [1, 2]: raise Exception("ssd_version can be only [1,2]") if location_loss_factor is not None: if not ((location_loss_factor > 0) and (location_loss_factor < 1)): raise Exception( '`location_loss_factor` should be greater than 0 and less than 1' ) self.location_loss_factor = location_loss_factor self._code = code self.ssd_version = ssd_version backbone_cut = None backbone_split = None if hasattr(self, '_orig_backbone'): self._backbone_ms = self._backbone self._backbone = self._orig_backbone _backbone_meta = cnn_config(self._orig_backbone) backbone_cut = _backbone_meta['cut'] backbone_split = _backbone_meta['split'] if backbone is None: self._backbone = models.resnet34 backbone_name = 'res' elif type(backbone) is str: self._backbone = getattr(models, backbone) backbone_name = backbone[:3] else: self._backbone = backbone backbone_name = 'custom' if not self._check_backbone_support(self._backbone): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) if self._backbone == models.mobilenet_v2: backbone_cut = -1 backbone_split = _mobilenet_split if ssd_version == 1: if grids == None: grids = [4, 2, 1] self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone, cut=backbone_cut), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] ssd_head = SSDHead(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) elif ssd_version == 2: # find bounding boxes height and width if grids is None: logger.info("Computing optimal grid size...") hw = data.height_width hw = np.array(hw) # find most suitable centroids for dataset centroid = kmeans(hw, 1) avg = avg_iou(hw, centroid) for num_anchor in range(2, 5): new_centroid = kmeans(hw, num_anchor) new_avg = avg_iou(hw, new_centroid) if (new_avg - avg) < 0.05: break avg = new_avg centroid = new_centroid.copy() # find grid size grids = list( map( int, map( round, data.chip_size / np.sort(np.max(centroid, axis=1))))) grids = list(set(grids)) grids.sort(reverse=True) if grids[-1] == 0: grids[-1] = 1 grids = list(set(grids)) self._create_anchors(grids, zooms, ratios) feature_sizes = model_sizes(create_body(self._backbone, cut=backbone_cut), size=(data.chip_size, data.chip_size)) num_features = feature_sizes[-1][-1] num_channels = feature_sizes[-1][1] if grids[0] > 8 and abs(num_features - grids[0] ) > 4 and backbone_name == 'res': num_features = feature_sizes[-2][-1] num_channels = feature_sizes[-2][1] backbone_cut = -3 ssd_head = SSDHeadv2(grids, self._anchors_per_cell, data.c, num_features=num_features, drop=drop, bias=bias, num_channels=num_channels) else: raise Exception('SSDVersion can only be 1 or 2') if hasattr(self, '_backbone_ms'): self._orig_backbone = self._backbone self._backbone = self._backbone_ms self.learn = cnn_learner(data=data, base_arch=self._backbone, cut=backbone_cut, split_on=backbone_split, custom_head=ssd_head) self._arcgis_init_callback() # make first conv weights learnable self.learn.model = self.learn.model.to(self._device) if focal_loss: self._loss_f = FocalLoss(data.c) else: self._loss_f = BCE_Loss(data.c) self.learn.loss_func = self._ssd_loss _set_multigpu_callback(self) if pretrained_path is not None: self.load(pretrained_path)
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
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)