[Explainer] PyTorch starter 0.857 F1-Score on public LB

from copy import copy

import numpy as np
from sklearn.model_selection import StratifiedKFold

import torch
torch.backends.cudnn.benchmark = True
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, Subset

import segmentation_models_pytorch as smp
import argus
from argus.callbacks import MonitorCheckpoint, EarlyStopping, LoggingToFile, ReduceLROnPlateau
import albumentations as A
from pytorch_toolbelt.inference.tiles import ImageSlicer, CudaTileMerger
from pytorch_toolbelt.losses import LovaszLoss

train_img = np.load('data_train.npz', allow_pickle=True, mmap_mode='r')['data']
train_labels = np.load('labels_train.npz', allow_pickle=True, mmap_mode='r')['labels']
test_img = np.load('data_test_1.npz', allow_pickle=True, mmap_mode='r')['data']

train_labels -= 1

train_img.shape, train_labels.shape, test_img.shape

adjusted_test_img = np.concatenate([train_img, test_img], axis=2)

_min, _max = adjusted_test_img.min(), adjusted_test_img.max()

train_img = (train_img - _min) / (_max - _min)
test_img = (test_img - _min) / (_max - _min)
adjusted_test_img = (adjusted_test_img - _min) / (_max - _min)

del adjusted_test_img

verticals = [0 for _ in range(train_img.shape[1])] + [1 for _ in range(train_img.shape[2])]
folds = list(StratifiedKFold(n_splits=5, random_state=42, shuffle=True).split(X=verticals, y=verticals))

class LovaszBCELoss(torch.nn.Module):
    def __init__(self, lovasz_weight=0.75, ce_weight=0.25):
        super().__init__()
        self.lovasz_weight = lovasz_weight
        self.ce_weight = ce_weight
        self.ce = torch.nn.CrossEntropyLoss()
        self.lovasz = LovaszLoss()

    def forward(self, output, target):
        if self.lovasz_weight > 0:
            lovasz = self.lovasz(torch.softmax(output, dim=1), target) * self.lovasz_weight
        else:
            lovasz = 0

        if self.ce_weight > 0:
            ce = self.ce(output, target.long()) * self.ce_weight
        else:
            ce = 0

        return lovasz + ce

class SeismicFaciesDataset(Dataset):
    def __init__(self, img, labels, train=True):
        self.img = img
        self.labels = labels
        self.xaxis = self.img.shape[1]
        self.yaxis = self.img.shape[2]

        self.aug = A.Compose([
            A.HorizontalFlip(p=0.5),
            A.ShiftScaleRotate(p=0.7, shift_limit=0, scale_limit=0.15, rotate_limit=15),
            A.RandomCrop(p=1, height=896, width=256),
        ])

    def __len__(self):
        return self.xaxis + self.yaxis

    def __getitem__(self, idx):
        if idx < self.xaxis:
            image, mask = self.img[:, idx], self.labels[:, idx]
        else:
            image, mask = self.img[:, :, idx-self.xaxis], self.labels[:, :, idx-self.xaxis]

        image = image[:, :, None]

        augmented = self.aug(image=image, mask=mask)
        image, mask = augmented['image'], augmented['mask']

        return image.transpose(2, 0, 1), mask

class SeismicFaciesModel(argus.Model):
    nn_module = smp.Unet
    optimizer = optim.SGD
    loss = LovaszBCELoss

params = {
    'nn_module': {
        'encoder_name': 'efficientnet-b3',
        'decoder_attention_type': 'scse',
        'classes': 6,
        'in_channels': 1,
        'activation': None
    },
    'loss': {
        'lovasz_weight': 0.75,
        'ce_weight': 0.25,
    },
    'optimizer': {'lr': 0.01, 'momentum': 0.9, 'weight_decay': 0.0001},
    'device': 'cuda'
}

def get_data_loaders(dataset, batch_size, train_index, test_index):
    train_dataset, test_dataset = Subset(dataset, train_index), Subset(copy(dataset), test_index)
    test_dataset.dataset.aug = A.PadIfNeeded(p=1, min_height=1024, min_width=800)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=16, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=16)
    return train_loader, test_loader

dataset = SeismicFaciesDataset(train_img, train_labels)

for i, (train_index, test_index) in enumerate(folds):
    model = SeismicFaciesModel(params)
    model.set_device((0, 1))

    train_loader, val_loader = get_data_loaders(dataset, batch_size=16, train_index=train_index, test_index=test_index)

    callbacks = [
        MonitorCheckpoint(dir_path=f'unet_fold_{i}', monitor='val_loss', max_saves=3),
        ReduceLROnPlateau(monitor='val_loss', patience=30, factor=0.64, min_lr=1e-8),
        EarlyStopping(monitor='val_loss', patience=50),
        LoggingToFile(f'unet_fold_{i}.log'),
    ]

    model.fit(train_loader,
          val_loader=val_loader,
          num_epochs=700,
          metrics=['loss'],
          callbacks=callbacks,
          metrics_on_train=False)

model = argus.load_model('unet_fold_0/model-692-0.053404.pth')

tiler = ImageSlicer(train_img.shape[:-1] + (1,), tile_size=(896, 256), tile_step=(1, 8))
merger = CudaTileMerger(tiler.target_shape, 6, tiler.weight)

test_labels = []
for img in test_img.transpose(2, 0, 1):
    tiles = [tile for tile in tiler.split(img[:, :, None])]

    for tiles_batch, coords_batch in DataLoader(list(zip(tiles, tiler.crops)), batch_size=96):
        tiles_batch = tiles_batch.permute(0, 3, 1, 2)
        pred_batch = torch.softmax(model.predict(tiles_batch), axis=1)
        merger.integrate_batch(pred_batch, coords_batch)

    merged_mask = merger.merge()
    merged_mask = merged_mask.permute(1, 2, 0).cpu().numpy()
    merged_mask = tiler.crop_to_orignal_size(merged_mask).argmax(2)
    
    test_labels.append(merged_mask)

test_labels = np.stack(test_labels).transpose(1, 2, 0)

np.savez_compressed(
    'prediction.npz',
    prediction=test_labels.astype(train_labels.dtype) + 1
)