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)