def train(arguments):
dls = UpscalerDataset(arguments.dataset_path,
arguments.resized_dataset_path)
feat_loss = get_loss(arguments.loss_model_path)
learner = unet_learner(dls.get_dataloaders(bs, size),
arch,
loss_func=feat_loss,
metrics=LossMetrics(feat_loss.metric_names),
blur=True,
norm_type=NormType.Weight)
# stage 1
print('stage 1')
do_fit(learner, wd, slice(lr * 10))
learner.unfreeze()
do_fit(learner, wd, slice(1e-5, lr))
# checkpoint
learner.save('checkpoint')
learner.load('checkpoint')
# stage 2
print('stage 2')
del learner.dls
learner.dls = dls.get_dataloaders(5, size * 2)
learner.freeze()
do_fit(learner, wd, slice(lr))
learner.unfreeze()
do_fit(learner, wd, slice(1e-6, 1e-4), pct_start=0.3)
# save
torch.save({
'model': learner.model,
'meta': make_meta(arguments)
}, arguments.output)
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, data, backbone=None, pretrained_path=None):
super().__init__()
self._device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
if not HAS_FASTAI:
_raise_fastai_import_error()
self._emd_template = _EMD_TEMPLATE
self._code = image_classifier_prf
if backbone is None:
self._backbone = models.resnet34
elif type(backbone) is str:
self._backbone = getattr(models, backbone)
else:
self._backbone = backbone
acc_metric = partial(accuracy, void_code=0, class_mapping=data.class_mapping)
self._data = data
self.learn = unet_learner(data, arch=self._backbone, metrics=acc_metric, wd=1e-2, bottle=True, last_cross=True)
self.learn.model = self.learn.model.to(self._device)
if pretrained_path is not None:
self.load(pretrained_path)
def __init__(self, data, backbone=None, pretrained_path=None, *args, **kwargs):
super().__init__(data, backbone)
feat_loss = create_loss(self._device.type)
data.c = 3
self.learn = unet_learner(data, arch=self._backbone, wd=1e-3, loss_func=feat_loss, callback_fns=LossMetrics, blur=True, norm_type=NormType.Weight)
self.learn.model = self.learn.model.to(self._device)
if pretrained_path is not None:
self.load(pretrained_path)
self._code = super_resolution
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 get_segmentation_learner(dls,
number_classes,
segmentation_type,
architecture_name,
backbone_name,
loss_func=None,
opt_func=Adam,
lr=defaults.lr,
splitter=trainable_params,
cbs=None,
pretrained=True,
normalize=True,
image_size=None,
metrics=None,
path=None,
model_dir='models',
wd=None,
wd_bn_bias=False,
train_bn=True,
moms=(0.95, 0.85, 0.95)):
"""This function return a learner for the provided architecture and backbone
Parameters:
dls (DataLoader): the dataloader to use with the learner
number_classes (int): the number of clases in the project. It should be >=2
segmentation_type (str): just `Semantic Segmentation` accepted for now
architecture_name (str): name of the architecture. The following ones are supported: `unet`, `deeplabv3+`, `hrnet`, `maskrcnn` and `u2^net`
backbone_name (str): name of the backbone
loss_func (): loss function.
opt_func (): opt function.
lr (): learning rates
splitter (): splitter function for freazing the learner
cbs (List[cb]): list of callbacks
pretrained (bool): it defines if a trained backbone is needed
normalize (bool):
image_size (int): REQUIRED for MaskRCNN. It indicates the desired size of the image.
metrics (List[metric]): list of metrics
path (): path parameter
model_dir (str): the path in which save models
wd (float): wieght decay
wd_bn_bias (bool):
train_bn (bool):
moms (Tuple(float)): tuple of different momentuns
Returns:
learner: value containing the learner object
"""
number_classes_name = ""
if number_classes == 2:
number_classes_name = "binary"
elif number_classes > 2:
number_classes_name = "multiple"
else:
raise Exception("The number of classes must be >=2")
check_architecture_configuration(number_classes=number_classes_name,
segmentation_type=segmentation_type,
architecture_name=architecture_name,
backbone_name=backbone_name)
learner = None
if architecture_name == "unet":
# TODO -> Revisar arch
learner = unet_learner(dls=dls,
arch=unet_backbone_name[backbone_name],
metrics=metrics,
wd=wd,
loss_func=loss_func,
opt_func=opt_func,
lr=lr,
splitter=splitter,
cbs=cbs,
path=path,
model_dir=model_dir,
wd_bn_bias=wd_bn_bias,
train_bn=train_bn,
pretrained=pretrained,
normalize=normalize,
moms=moms)
elif architecture_name == "deeplabv3+":
model = DeepLabV3Plus(backbone_name=backbone_name,
nclass=number_classes,
pretrained=pretrained)
learner = Learner(dls=dls,
model=model,
loss_func=loss_func,
opt_func=opt_func,
lr=lr,
splitter=splitter,
cbs=cbs,
metrics=metrics,
path=path,
model_dir=model_dir,
wd=wd,
wd_bn_bias=wd_bn_bias,
train_bn=train_bn)
elif architecture_name == "hrnet":
model = HRNet(nclass=number_classes,
backbone_name=backbone_name,
pretrained=pretrained)
learner = Learner(dls=dls,
model=model,
loss_func=loss_func,
opt_func=opt_func,
lr=lr,
splitter=splitter,
cbs=cbs,
metrics=metrics,
path=path,
model_dir=model_dir,
wd=wd,
wd_bn_bias=wd_bn_bias,
train_bn=train_bn)
elif architecture_name == "maskrcnn":
if image_size is None:
raise Exception(
"MaskRCNN need to define image_size. This values are for reescaling the image"
)
model = maskrcnn_resnet50_fpn(num_classes=number_classes,
min_size=image_size,
max_size=image_size)
learner = mask_rcnn.MaskRCNNLearner(dls=dls,
model=model,
loss_func=loss_func,
opt_func=opt_func,
lr=lr,
splitter=splitter,
cbs=cbs,
metrics=metrics,
path=path,
model_dir=model_dir,
wd=wd,
wd_bn_bias=wd_bn_bias,
train_bn=train_bn)
elif architecture_name == "u2^net":
model = None
if backbone_name == "small":
model = u2net.U2NETP(3, 1)
elif backbone_name == "normal":
model = u2net.U2NET(3, 1)
learner = u2net.USquaredNetLearner(dls=dls,
model=model,
opt_func=opt_func,
lr=lr,
splitter=splitter,
cbs=cbs,
metrics=metrics,
path=path,
model_dir=model_dir,
wd=wd,
wd_bn_bias=wd_bn_bias,
train_bn=train_bn)
return learner
def run():
models = {
'resnet34': mod.resnet34,
'resnet50': mod.resnet50,
'resnet101': mod.resnet101,
'resnet152': mod.resnet152
}
save_name = f'seg_{cfg.MODEL}_{cfg.TRAIN_SIZE}'
save_name = f'{save_name}_{getNextFilePath(cfg.MODELS_PATH, save_name)}'
test_list = (SegmentationItemList.from_folder(cfg.TEST_PATH,
extensions=['.dcm']))
best = 0
pred_path = cfg.PRED_PATH / save_name
if not pred_path.is_dir():
pred_path.mkdir()
project = neptune.init('schwobr/SIIM-Pneumothorax')
for k, db in enumerate(
load_data_kfold(cfg.LABELS,
bs=cfg.BATCH_SIZE,
seed=cfg.SEED,
train_size=cfg.TRAIN_SIZE,
xtra_tfms=[gaussian_noise()])):
print(f'fold {k}')
learner = unet_learner(db,
models[cfg.MODEL],
pretrained=cfg.PRETRAINED,
wd=cfg.WD,
model_dir=cfg.MODELS_PATH,
metrics=[dice],
loss_func=URLoss(CrossEntropyLoss(),
reduction='sum',
do_spatial_reduc=True))
fold_name = f'fold{k}_' + save_name
set_BN_momentum(learner.model, momentum=0.01)
# replace_bn(learner, learner.model)
learner.fit_one_cycle(
cfg.EPOCHS,
slice(cfg.LR),
callbacks=[
SaveModelCallback(learner, monitor='dice', name=fold_name),
AccumulateScheduler(learner, min(8, 64 // cfg.BATCH_SIZE)),
# LookAhead(learner, k=max(6, 64 // cfg.BATCH_SIZE)),
NeptuneCallback(learner,
project,
name=fold_name,
params={
'lr': cfg.LR,
'wd': cfg.WD,
'size': cfg.TRAIN_SIZE
})
])
met = max([met[0] for met in learner.recorder.metrics])
if met > best:
learner.save(save_name)
best = met
print(f'New best fold {k} with dice {best}')
# learner.neptune_callback.send_artifact(
# cfg.MODELS_PATH/(fold_name+'.pth'))
learner.neptune_callback.stop()
learner.unfreeze()
fold_name = 'uf_' + fold_name
learner.fit_one_cycle(cfg.UNFROZE_EPOCHS,
slice(cfg.UF_LR / 100, cfg.UF_LR),
callbacks=[
SaveModelCallback(learner,
monitor='dice',
name=fold_name),
AccumulateScheduler(learner,
64 // cfg.BATCH_SIZE),
NeptuneCallback(learner,
project,
name=fold_name,
params={
'lr': cfg.LR,
'wd': cfg.WD,
'size': cfg.TRAIN_SIZE
})
])
met = max([met[0] for met in learner.recorder.metrics])
if met > best:
learner.save(save_name)
best = met
print(f'New best fold {k} with dice {best}')
# learner.neptune_callback.send_artifact(
# cfg.MODELS_PATH/(fold_name+'.pth'))
learner.neptune_callback.stop()
learner.data.add_test(test_list,
label=[test_list.items[0], '-1'],
tfms=(),
tfm_y=True)
save_preds(learner, pred_path / str(k))
exp = project.create_experiment(name=save_name,
description='k-fold mtl training',
params={
'lr': cfg.LR,
'wd': cfg.WD,
'size': cfg.TRAIN_SIZE
})
# exp.send_artifact(cfg.MODELS_PATH/(save_name+'.pth'))
learner.load(save_name)
learner.data.add_test(test_list,
label=[test_list.items[0], '-1'],
tfms=(),
tfm_y=True)
thr = get_best_thr(learner,
plot=False,
test_size=cfg.TEST_SIZE,
exp=None,
fig_path=cfg.FIG_PATH / (save_name + '.png'))
create_submission_kfold(learner,
cfg.SUB_PATH / (save_name + '.csv'),
pred_path,
test_size=cfg.TEST_SIZE,
thr=thr)
exp.send_artifact(cfg.SUB_PATH / (save_name + '.csv'))
exp.stop()
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)
