Exemplo n.º 1
0
Arquivo: unet.py Projeto: zyzju/PSSR
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
Exemplo n.º 2
0
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
Exemplo n.º 3
0
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
Exemplo n.º 4
0
    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)
Exemplo n.º 5
0
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
Exemplo n.º 6
0
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
Exemplo n.º 7
0
    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)
Exemplo n.º 8
0
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
Exemplo n.º 10
0
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
Exemplo n.º 12
0
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
Exemplo n.º 13
0
    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)
Exemplo n.º 14
0
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
Exemplo n.º 15
0
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
Exemplo n.º 16
0
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
Exemplo n.º 17
0
Arquivo: model.py Projeto: iskode/acdc
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
Exemplo n.º 18
0
    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)
Exemplo n.º 20
0
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
Exemplo n.º 21
0
                         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)
Exemplo n.º 22
0
    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)