def wnres_unet_model(in_c, out_c, arch, blur=True, blur_final=True, self_attention=True, last_cross=True, bottle=True, norm_type=NormType.Weight, **wnres_args): meta = cnn_config(arch) enc_model = arch(c_in=in_c) cut = cnn_config(arch)['cut'] body = nn.Sequential(*list(enc_model.children())[:cut]) model = DynamicUnet(body, n_classes=out_c, blur=blur, blur_final=blur_final, self_attention=self_attention, norm_type=norm_type, last_cross=last_cross, bottle=bottle, **wnres_args) return model, meta
def unet_learner(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = NormType, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, cut: Union[int, Callable] = None, hypercolumns=True, **learn_kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = cnn_config(arch) body = create_body(arch, pretrained, cut) M = DynamicUnet_Hcolumns if hypercolumns else DynamicUnet model = to_device( M(body, n_classes=data.c, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle), data.device) learn = Learner(data, model, **learn_kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
def iqa_cnn_learner(data: DataBunch, base_arch: Callable, cut: Union[int, Callable] = None, pretrained: bool = True, lin_ftrs: Optional[Collection[int]] = None, ps: Floats = 0.5, custom_head: Optional[nn.Module] = None, split_on: Optional[SplitFuncOrIdxList] = None, bn_final: bool = False, init=nn.init.kaiming_normal_, concat_pool: bool = True, learner=IqaLearner, **kwargs: Any): "Build convnet style learner." from fastai.vision.learner import cnn_config meta = cnn_config(base_arch) model = create_cnn_model(base_arch, data.c, cut, pretrained, lin_ftrs, ps=ps, custom_head=custom_head, bn_final=bn_final, concat_pool=concat_pool) model.__name__ = base_arch.__name__ learn = learner(data, model, **kwargs) learn.split(split_on or meta['split']) if pretrained: learn.freeze() if init: apply_init(model[1], init) return learn
def __init__(self, base_arch, no_diseases, dropout=0.5, init=nn.init.kaiming_normal_): super(CNNPretrainedModel, self).__init__() self.body = create_body(base_arch, pretrained=True) nf = num_features_model(nn.Sequential(*self.body.children())) * 2 self.disease_head = create_head(nf, no_diseases, ps=0.5, concat_pool=True, bn_final=False) #self.age_head = create_head(nf, 1, ps=0.5, concat_pool=True, bn_final=False) self.gender_head = create_head(nf, 2, ps=0.5, concat_pool=True, bn_final=False) #self.projection_head = create_head(nf, 3, ps=0.5, concat_pool=True, bn_final=False) self.disease_model = nn.Sequential(self.body, self.disease_head) self.meta = cnn_config(base_arch) self.split(self.meta['split']) self.freeze() apply_init(self.disease_head, init) #apply_init(self.age_head, init) apply_init(self.gender_head, init)
def create_cnn(data, arch, pretrained=False, is_mono_input=True, **kwargs): meta = cnn_config(arch) body = create_body(arch, pretrained) # sum up the weights of in_channels axis, to reduce to single input channel # Suggestion by David Gutman # https://forums.fast.ai/t/black-and-white-images-on-vgg16/2479/2 if is_mono_input: first_conv_layer = body[0] first_conv_weights = first_conv_layer.state_dict()['weight'] assert first_conv_weights.size(1) == 3 # RGB channels dim summed_weights = torch.sum(first_conv_weights, dim=1, keepdim=True) first_conv_layer.weight.data = summed_weights first_conv_layer.in_channels = 1 else: # In this case, the input is a stereo first_conv_layer = body[0] first_conv_weights = first_conv_layer.state_dict()['weight'] assert first_conv_weights.size(1) == 3 # RGB channels dim summed_weights = torch.sum(first_conv_weights, dim=1, keepdim=True) first_conv_layer.weight.data = first_conv_weights[:, : 2, :, :] # Keep only 2 channels for the weights first_conv_layer.in_channels = 2 nf = num_features_model(body) * 2 head = create_head(nf, data.c, None, 0.5) model = nn.Sequential(body, head) learn = Learner(data, model, **kwargs) learn.split(meta['split']) if pretrained: learn.freeze() apply_init(model[1], nn.init.kaiming_normal_) return learn
def Learner_create_unet(cls, data:DataBunch, arch:Callable, pretrained:bool=True, split_on:Optional[SplitFuncOrIdxList]=None, **kwargs:Any)->None: "Build Unet learners." meta = cnn_config(arch) body = create_body(arch(pretrained), meta['cut']) model = to_device(models.unet.DynamicUnet(body, n_classes=data.c), data.device) learn = cls(data, model, **kwargs) learn.split(ifnone(split_on,meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
def __init__(self, base_arch, no_classes, dropout=0.5, init=nn.init.kaiming_normal_): super(CNNPretrainedModel, self).__init__() self.model = create_cnn_model(base_arch, no_classes, ps=dropout) self.meta = cnn_config(base_arch) self.split(self.meta['split']) self.freeze() apply_init(self.model[1], init)
def siamese_learner(data, pretrained_model_class, emsize=128, margin=1.0, callback_fns=None): meta = cnn_config(pretrained_model_class) model = create_cnn_model(pretrained_model_class, emsize) model = SiameseNet(model) learn = Learner(data, model, loss_func=ContrastiveLoss(margin), callback_fns=callback_fns) learn.split(meta["split"](model.embedding_net)) apply_init(model.embedding_net[1], nn.init.kaiming_normal_) return learn
def create_cnn(data:DataBunch, arch:Callable, cut:Union[int,Callable]=None, pretrained:bool=True, lin_ftrs:Optional[Collection[int]]=None, ps:Floats=0.5, custom_head:Optional[nn.Module]=None, split_on:Optional[SplitFuncOrIdxList]=None, bn_final:bool=False, **kwargs:Any)->Learner: "Build convnet style learners." meta = cnn_config(arch) body = create_body(arch(pretrained), ifnone(cut,meta['cut'])) # nf = num_features_model(body) * 2 head = custom_head or create_head(2048, data.c, lin_ftrs, ps=ps, bn_final=bn_final) model = nn.Sequential(body, head) learn = Learner(data, model, **kwargs) learn.split(ifnone(split_on,meta['split'])) if pretrained: learn.freeze() apply_init(model[1], nn.init.kaiming_normal_) return learn
def siamese_embedding_learner(data, pretrained_model_class, emsize=128, margin=1.0, callback_fns=None): meta = cnn_config(pretrained_model_class) model = create_cnn_model(pretrained_model_class, emsize) learn = Learner( data, model, loss_func=OnlineContrastiveLoss(margin, HardNegativePairSelector()), callback_fns=callback_fns, ) learn.split(meta["split"]) apply_init(model[1], nn.init.kaiming_normal_) return learn
def create_cnn(data:DataBunch, arch:Callable, cut:Union[int,Callable]=None, pretrained:bool=True, lin_ftrs:Optional[Collection[int]]=None, ps:Floats=0.5, custom_head:Optional[nn.Module]=None, split_on:Optional[SplitFuncOrIdxList]=None, classification:bool=True, **kwargs:Any)->Learner: "Build convnet style learners." assert classification, 'Regression CNN not implemented yet, bug us on the forums if you want this!' meta = cnn_config(arch) body = create_body(arch(pretrained), ifnone(cut,meta['cut'])) head = custom_head model = nn.Sequential(body, head) learner_cls = ifnone(data.learner_type(), ClassificationLearner) learn = learner_cls(data, model, **kwargs) learn.split(ifnone(split_on,meta['split'])) if pretrained: learn.freeze() apply_init(model[1], nn.init.kaiming_normal_) return learn
def create_body_local(arch: Callable, pretrained: bool = True, cut: Optional[Union[int, Callable]] = None): "Cut off the body of a typically pretrained `model` at `cut` (int) or cut the model as specified by `cut(model)` (function)." if pretrained: model = arch(num_classes=1000, pretrained='imagenet') cut = ifnone(cut, cnn_config(arch)['cut']) if cut is None: ll = list(enumerate(model.children())) cut = next(i for i, o in reversed(ll) if has_pool_type(o)) if isinstance(cut, int): return nn.Sequential(*(list(model.children())[:cut])) elif isinstance(cut, Callable): return cut(model) else: pass
def __init__(self, data, backbone=None, pretrained_path=None): super().__init__(data, backbone) self._code = image_classifier_prf backbone_cut = None backbone_split = None _backbone = self._backbone if hasattr(self, '_backbone_'): _backbone = self._backbone_ if not (self._check_backbone_support(_backbone)): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) if hasattr(self, '_backbone_'): _backbone_meta = cnn_config(self._backbone_) backbone_cut = _backbone_meta['cut'] backbone_split = _backbone_meta['split'] acc_metric = partial(accuracy, void_code=0, class_mapping=data.class_mapping) self.learn = unet_learner(data, arch=self._backbone, metrics=acc_metric, wd=1e-2, bottle=True, last_cross=True, cut=backbone_cut, split_on=backbone_split) self._arcgis_init_callback() # make first conv weights learnable self.learn.callbacks.append(LabelCallback( self.learn)) #appending label callback self.learn.model = self.learn.model.to(self._device) # _set_multigpu_callback(self) # MultiGPU doesn't work for U-Net. (Fastai-Forums) if pretrained_path is not None: self.load(pretrained_path)
def DR_learner(data: DataBunch, base_arch: Callable, cut: Union[int, Callable] = None, pretrained: bool = True, lin_ftrs: Optional[Collection[int]] = None, ps: Floats = 0.5, custom_head: Optional[nn.Module] = None, split_on: Optional[SplitFuncOrIdxList] = None, bn_final: bool = False, init=nn.init.kaiming_normal_, concat_pool: bool = True, attention=False, **kwargs: Any) -> Learner: "Build convnet style learner." meta = cnn_config(base_arch) "Create custom convnet architecture" body = create_body(base_arch, pretrained, cut) if custom_head is None: # quadnet head nf = num_features_model( nn.Sequential(*body.children())) * (2 if concat_pool else 1) if attention: pass head = create_DR_head(nf, data.c + 2, lin_ftrs, ps=ps, concat_pool=concat_pool, bn_final=bn_final) else: head = custom_head model = nn.Sequential(body, head) learn = Learner(data, model, **kwargs) learn.split(split_on or meta['split']) if pretrained: learn.freeze() if init: apply_init(model[1], init) return learn
def unet_learner_wide_exp(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = NormType, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, nf_factor: int = 1, nf: int = 128, **kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = cnn_config(arch) body = create_body(arch, pretrained) model = to_device( DynamicUnetWide( body, n_classes=data.c, #for MSGAN.was data.c blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle, nf_factor=nf_factor, ), data.device, ) learn = Learner(data, model, **kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
def unet_learner_wide(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = NormType, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, nf_factor: int = 1, **kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = cnn_config(arch) body = create_body(arch, pretrained) #preModel = torch.load('/Projects/mk_utils/Convert_Mxnet_to_Pytorch/Pytorch_NewModel.pth') #preModel = list(preModel.children())[0] #body = pthutils.cut_model(preModel,-3) model = to_device( DynamicUnetWide(body, n_classes=data.c, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle, nf_factor=nf_factor), data.device) #print(model[2]) learn = Learner(data, model, **kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze_to(-2) #print(model[2]) #apply_init(model[2], nn.init.kaiming_normal_) return learn
def unet4mri_learner(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = None, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, cut: Union[int, Callable] = None, **learn_kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = ln.cnn_config(arch) body = ln.create_body(arch, pretrained, cut) try: size = data.train_ds[0][0].size except: size = next(iter(data.train_dl))[0].shape[-2:] model = to_device( MyDynamicUnet(body, bs=data.batch_size, n_classes=data.c, img_size=size, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle), data.device) learn = Learner(data, model, **learn_kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() apply_init(model[2], nn.init.kaiming_normal_) return learn
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)
every='improvement', name=Path(MODEL_DIR.name) / f'fold{fold_num}'), partial(CSVLogger, filename=LOG_FNAME) ] learn = cnn_learner(data=fold_data, base_arch=arch, metrics=[accuracy, auc], callbacks=learn_callbacks, callback_fns=learn_callback_fns) # change model to 4D input convert_model_to_4D_input(learn.model) # re-assign new 4d parameters to layer groups and freeze model meta = cnn_config(arch) learn.split(meta['split']) learn.freeze() # Stage-1 training learn.lr_find() try: learn.recorder.plot(suggestion=True, k=5) except: learn.recorder.plot(suggestion=True) max_lr = learn.recorder.min_grad_lr print(f"Stage-1 training with lr={max_lr}") learn.fit_one_cycle(epochs['stage1'], max_lr=max_lr) # Stage-2 training learn.freeze_to(1)
def unet_learner_deep_MS(data: DataBunch, arch: Callable, pretrained: bool = True, blur_final: bool = True, norm_type: Optional[NormType] = NormType, split_on: Optional[SplitFuncOrIdxList] = None, blur: bool = False, self_attention: bool = False, y_range: Optional[Tuple[float, float]] = None, last_cross: bool = True, bottle: bool = False, nf_factor: float = 1.5, nc=3, nz=8, sz=128, **kwargs: Any) -> Learner: "Build Unet learner from `data` and `arch`." meta = cnn_config(arch) body = create_body(arch, pretrained) model = to_device( DynamicUnetDeep( body, n_classes=data.c, blur=blur, blur_final=blur_final, self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross, bottle=bottle, nf_factor=nf_factor, ), data.device, ) learn = Learner(data, model, **kwargs) learn.split(ifnone(split_on, meta['split'])) if pretrained: learn.freeze() # adjust layers to allow 11D input # could be appended with a new network at position 0 instead, but that way it would not fit an 11GB GPU # thus, replace the first layer with 11D input learn.model[0][0] = nn.Conv2d((nc + nz), int(sz / 2), kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) learn.layer_groups[0][0] = nn.Conv2d((nc + nz), int(sz / 2), kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) # and replace the merge layer learn.model[10][0][0] = nn.Conv2d(311, 311, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) learn.layer_groups[-1][62] = nn.Conv2d(311, 311, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) learn.model[10][1][0] = nn.Conv2d(311, 311, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) learn.layer_groups[-1][64] = nn.Conv2d(311, 311, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) learn.model[11][0] = nn.Conv2d(311, 3, kernel_size=(1, 1), stride=(1, 1)) learn.layer_groups[-1][67] = nn.Conv2d(311, 3, kernel_size=(1, 1), stride=(1, 1)) # initialize the newly created layers apply_init(model[0], nn.init.kaiming_normal_) apply_init(model[10][0][0], nn.init.kaiming_normal_) apply_init(model[10][1][0], nn.init.kaiming_normal_) apply_init(model[11][0], nn.init.kaiming_normal_) apply_init(model[2], nn.init.kaiming_normal_) return learn
num_workers=NUM_WORKERS, normalization=STATS) # ============================== model learning ============================== model = create_cnn_model(base_arch=ARCH, nc=NUM_CLASSES) model.load_state_dict(torch.load(STATE_FILE)['model']) model = to_device(model, DEVICE) learn = AccGradLearner(data, model, metrics=[METRIC], model_dir=LOGGING_FOLDER) learn.split(cnn_config(model)['split']) set_BN_momentum(learn.model, batch_size=BATCH_SIZE) learn.clip_grad(1.) # callbacks csv_logger = CSVLogger(learn=learn, filename=f'{LOGGING_FOLDER}/fit_trace', append=True) early_stopping = EarlyStoppingCallback(learn=learn, monitor='accuracy', patience=PATIENCE) save_model = SaveModelCallback(learn=learn, monitor='accuracy', name='best_model') acc_grad = AccumulateStep(learn, 64 // BATCH_SIZE)
def __init__(self, data, backbone=None, pretrained_path=None, backend='pytorch', *args, **kwargs): self._backend = backend if self._backend == 'tensorflow': super().__init__(data, None) self._intialize_tensorflow(data, backbone, pretrained_path, kwargs) else: super().__init__(data, backbone) # import pdb; pdb.set_trace(); self._ignore_classes = kwargs.get('ignore_classes', []) if self._ignore_classes != [] and len(data.classes) <= 3: raise Exception( f"`ignore_classes` parameter can only be used when the dataset has more than 2 classes." ) data_classes = list(self._data.class_mapping.keys()) if 0 not in list(data.class_mapping.values()): self._ignore_mapped_class = [ data_classes.index(k) + 1 for k in self._ignore_classes if k != 0 ] else: self._ignore_mapped_class = [ data_classes.index(k) + 1 for k in self._ignore_classes ] if self._ignore_classes != []: if 0 not in self._ignore_mapped_class: self._ignore_mapped_class.insert(0, 0) global accuracy accuracy = partial( accuracy, ignore_mapped_class=self._ignore_mapped_class) self.mixup = kwargs.get('mixup', False) self.class_balancing = kwargs.get('class_balancing', False) self.focal_loss = kwargs.get('focal_loss', False) self._code = image_classifier_prf backbone_cut = None backbone_split = None _backbone = self._backbone if hasattr(self, '_orig_backbone'): _backbone = self._orig_backbone if not (self._check_backbone_support(_backbone)): raise Exception( f"Enter only compatible backbones from {', '.join(self.supported_backbones)}" ) if hasattr(self, '_orig_backbone'): _backbone_meta = cnn_config(self._orig_backbone) backbone_cut = _backbone_meta['cut'] backbone_split = _backbone_meta['split'] if not _isnotebook() and arcgis_os.name == 'posix': _set_ddp_multigpu(self) if self._multigpu_training: self.learn = unet_learner( data, arch=self._backbone, metrics=accuracy, wd=1e-2, bottle=True, last_cross=True, cut=backbone_cut, split_on=backbone_split).to_distributed( self._rank_distributed) else: self.learn = unet_learner(data, arch=self._backbone, metrics=accuracy, wd=1e-2, bottle=True, last_cross=True, cut=backbone_cut, split_on=backbone_split) else: self.learn = unet_learner(data, arch=self._backbone, metrics=accuracy, wd=1e-2, bottle=True, last_cross=True, cut=backbone_cut, split_on=backbone_split) if self.class_balancing: if data.class_weight is not None: class_weight = torch.tensor( [data.class_weight.mean()] + data.class_weight.tolist()).float().to(self._device) else: if getattr(data, 'overflow_encountered', False): logger.warning( "Overflow Encountered. Ignoring `class_balancing` parameter." ) class_weight = [1] * len(data.classes) else: logger.warning( "Could not find 'NumPixelsPerClass' in 'esri_accumulated_stats.json'. Ignoring `class_balancing` parameter." ) if self._ignore_classes != []: if not self.class_balancing: class_weight = torch.tensor([1] * data.c).float().to( self._device) class_weight[self._ignore_mapped_class] = 0. else: class_weight = None self.learn.loss_func = CrossEntropyFlat(class_weight, axis=1) if self.focal_loss: self.learn.loss_func = FocalLoss(self.learn.loss_func) if self.mixup: self.learn.callbacks.append(MixUpCallback(self.learn)) self._arcgis_init_callback() # make first conv weights learnable self.learn.callbacks.append(LabelCallback( self.learn)) #appending label callback self.learn.model = self.learn.model.to(self._device) # _set_multigpu_callback(self) # MultiGPU doesn't work for U-Net. (Fastai-Forums) if pretrained_path is not None: self.load(pretrained_path)