<jupyter_start><jupyter_text>Segmentation training Unet
> Training notebook with Tensorboard logs.  Setup<jupyter_code># hide
%load_ext autoreload
%autoreload 2
%matplotlib inline
from fastai.vision.all import *
from fastai.callback.tensorboard import TensorBoardCallback

import random
import cv2
import pandas as pd
import numpy as np

from matplotlib import pyplot as plt

import torch
import albumentations as alb

import segmentation_models_pytorch as smp

from steel_segmentation.utils import get_train_df
from steel_segmentation.transforms import SteelDataBlock, SteelDataLoaders
from steel_segmentation.losses import MultiClassesSoftBCEDiceLoss, LossEnabler
from steel_segmentation.metrics import ModDiceMulti
from steel_segmentation.optimizer import opt_func
def seed_everything(seed=69):
    """
    Seeds `random`, `os.environ["PYTHONHASHSEED"]`,
    `numpy`, `torch.cuda` and `torch.backends`.
    """
    #warnings.filterwarnings("ignore")
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_everything()<jupyter_output><empty_output><jupyter_text>Training parameters:<jupyter_code>bs = 16
size = (224,512)
epochs = 30
lr = 3e-4
path = Path("../data") # where data dir is<jupyter_output><empty_output><jupyter_text>Data loading<jupyter_code>df = get_train_df(path, only_faulty=True, pivot=True)
df.describe(include="all")
def get_train_aug(height, width): 
    tfm_list = [
        alb.RandomCrop(height, width, p=1.0),
        alb.OneOf(
          [
           alb.VerticalFlip(p=0.5),
           alb.HorizontalFlip(p=0.5),
          ], p=0.5),
        alb.RandomBrightnessContrast(
            brightness_limit=(-0.2, 0.2), contrast_limit=(-0.2, 0.2), p=0.5),
    ]
    return alb.Compose(tfm_list)

def get_valid_aug(height, width): 
    tfms = [alb.RandomCrop(height, width, p=1.0)]
    return alb.Compose(tfms)

device = "cuda" if torch.cuda.is_available() else "cpu"
train_aug = get_train_aug(*size)
valid_aug = get_valid_aug(*size)
block = SteelDataBlock(path, train_aug=train_aug, valid_aug=valid_aug)
dls = SteelDataLoaders(block, df, bs=bs, device=device)
xb, yb = dls.one_batch()
print(xb.shape, yb.shape)<jupyter_output>torch.Size([16, 3, 224, 512]) torch.Size([16, 4, 224, 512])<jupyter_text>Model<jupyter_code>model = smp.Unet(encoder_name="resnet18", encoder_weights="imagenet", classes=4, activation=None)
criterion = BCEWithLogitsLossFlat(pos_weight=torch.tensor([2.0,2.0,1.0,1.5])) # pos_weight because class imbalance
#opt_func = partial(opt_func, torch_opt=torch.optim.Adam) # no need to use pytorch optim
opt_func = RAdam
model_dir = Path("../models")
metrics = [ModDiceMulti(with_logits=True)]
learner = Learner(
    dls = dls,
    model = model,
    loss_func = criterion,
    opt_func = opt_func,
    metrics = metrics,
    model_dir = model_dir,
    cbs = [LossEnabler]
)
learner.summary()
learner.show_training_loop()<jupyter_output>Start Fit
   - before_fit     : [TrainEvalCallback, Recorder, ProgressCallback]
  Start Epoch Loop
     - before_epoch   : [Recorder, ProgressCallback]
    Start Train
       - before_train   : [TrainEvalCallback, Recorder, ProgressCallback]
      Start Batch Loop
         - before_batch   : []
         - after_pred     : [LossEnabler]
         - after_loss     : []
         - before_backward: []
         - before_step    : []
         - after_step     : []
         - after_cancel_batch: []
         - after_batch    : [TrainEvalCallback, Recorder, ProgressCallback]
      End Batch Loop
    End Train
     - after_cancel_train: [Recorder]
     - after_train    : [Recorder, ProgressCallback]
    Start Valid
       - before_validate: [TrainEvalCallback, Recorder, ProgressCallback]
      Start Batch Loop
         - **CBs same as train batch**: []
      End Batch Loop
    End Valid
     - after_cancel_validate: [Recorder]
     - after_validate : [Recorder, ProgressCallback]
  End Epoch Loop
[...]<jupyter_text>Logging with the TensorBoardCallback:<jupyter_code># logging info
log_dir = Path("../logs") / f"unet_resnet_bce_epochs{epochs}_lr{lr}"
log_dir
train_cbs = [
    TensorBoardCallback(log_dir=log_dir, log_preds=True, trace_model=True, projector=False),
    GradientAccumulation(n_acc=24),
    SaveModelCallback(monitor="valid_loss", fname=log_dir.name, with_opt=True),
]
learner.fit(epochs, lr=lr, cbs=train_cbs)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1113561190738699, 'fraction_numerical': 0.014773980154355016, 'mean_word_length': 3.7898627243928193, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '174746', 'license': 'Apache-2.0', 'n_tokens_mistral': 1658, 'n_tokens_neox': 1588, 'n_words': 303, 'path': 'nbs/10_segmentation_unet_resnet.ipynb', 'repo_name': 'marcomatteo/steel-segmentation-nbdev', 'size': '4535'}