import pandas as pd
import numpy as np
import os
from sklearn.model_selection import GroupKFold
from sklearn.metrics import roc_auc_score, roc_curve, log_loss
import torch
from torchvision import models, transforms
import torch.nn.functional as F
from torch.utils.data import DataLoader
import time
import apex
import glob
from shutil import copyfile
from datetime import datetime
from colorama import init, Fore, Back, Style
from collections import OrderedDict
import matplotlib.pyplot as plt
import sys
sys.path.append("/home/anjum/PycharmProjects/kaggle")
# sys.path.append("/home/anjum/rsna_code") # GCP
from rsna_intracranial_hemorrhage_detection.datasets import ICHDataset, BalancedRandomSampler
from rsna_intracranial_hemorrhage_detection.utils import *
from rsna_intracranial_hemorrhage_detection.networks import model_builder
torch.multiprocessing.set_sharing_strategy('file_system')
init(autoreset=True)
INPUT_DIR = "/mnt/storage_dimm2/kaggle_data/rsna-intracranial-hemorrhage-detection/"
OUTPUT_DIR = "/mnt/storage/kaggle_output/rsna-intracranial-hemorrhage-detection/"
# INPUT_DIR = "/home/anjum/rsna_data/" # GCP
# OUTPUT_DIR = "/home/anjum/rsna_output/" # GCP
N_WORKERS = 4
FOLDS = 5
def competition_metric(target, predictions):
target = target.numpy()
predictions = torch.sigmoid(predictions).numpy()
return log_loss(target.flatten(), predictions.flatten(), sample_weight=[2, 1, 1, 1, 1, 1] * target.shape[0])
# See https://www.kaggle.com/c/rsna-intracranial-hemorrhage-detection/discussion/109995#636452
def weighted_multi_label_logloss(criterion, prediction, target, weights):
assert target.size() == prediction.size()
assert weights.shape[0] == target.size(1)
loss = 0
for i in range(target.size(1)):
loss += weights[i] * criterion(prediction[:, i], target[:, i])
return loss
def train_model(args, device="cpu"):
start = time.time()
metadata = pd.read_parquet(os.path.join(INPUT_DIR, 'train_metadata.parquet.gzip'))
# metadata = metadata.sample(1024*10, random_state=48) # For testing
train_transforms = transforms.Compose([transforms.ToPILImage(),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(degrees=15),
transforms.RandomResizedCrop(size=(args.img_size, args.img_size),
scale=(0.85, 1.0), ratio=(0.8, 1.2)),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406],
# std=[0.229, 0.224, 0.225])
])
test_loaders = build_tta_loaders(args.img_size, dataset=args.stage, batch_size=args.batch, num_workers=N_WORKERS,
image_folder=args.img_folder)
# out-of-fold predictions on train data and averaged predictions on test data
n_images = metadata.shape[0]
categories = ["any", "epidural", "intraparenchymal", "intraventricular", "subarachnoid", "subdural"]
oof_df = pd.DataFrame({c: np.zeros(n_images) for c in categories}, index=metadata['SOPInstanceUID'])
prediction = np.zeros((len(test_loaders[0].dataset), len(categories)))
# list of scores on folds
scores, fold_data = [], OrderedDict()
# kfold = GroupKFold(n_splits=FOLDS)
# for fold_n, (train_index, valid_index) in enumerate(kfold.split(metadata, groups=metadata["PatientID"])):
# train_image_ids = metadata.iloc[train_index]['SOPInstanceUID'].values
# valid_image_ids = metadata.iloc[valid_index]['SOPInstanceUID'].values
# Using the same fold definition as the rest of the team
for fold_n in range(5):
train_image_ids = pd.read_csv(os.path.join(INPUT_DIR, args.cv_scheme, f"train_{fold_n}.csv"))["sop_instance_uid"]
valid_image_ids = pd.read_csv(os.path.join(INPUT_DIR, args.cv_scheme, f"valid_{fold_n}.csv"))["sop_instance_uid"]
train_dataset = ICHDataset("train", phase=0, image_filter=train_image_ids, transforms=train_transforms,
image_folder=args.img_folder)
train_loader = DataLoader(train_dataset, batch_size=args.batch, num_workers=N_WORKERS, shuffle=True,
drop_last=True)
valid_loaders = build_tta_loaders(args.img_size, dataset="train", batch_size=args.batch,
phase=0, image_filter=valid_image_ids, num_workers=N_WORKERS,
image_folder=args.img_folder)
# Create fresh network. optimiser etc.
torch.manual_seed(48 + fold_n)
model = model_builder(args.architecture)
model = model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5,
# patience=5, min_lr=0.00001, verbose=True)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
checkpoint_file = os.path.join(args.output_dir, f"phase0_fold_{fold_n + 1}.pt")
early_stopping = EarlyStopping(patience=3, verbose=True, file_path=checkpoint_file, parallel=args.parallel)
train_loss, valid_loss, valid_metric = [], [], []
weights = np.array([2, 1, 1, 1, 1, 1])
weights_norm = weights / weights.sum()
criterion = torch.nn.BCEWithLogitsLoss()
# Resume training
if args.checkpoint_path is not None:
try:
checkpoint_file = max(glob.glob(args.checkpoint_path + f'/*fold_{fold_n + 1}*.pt'),
key=os.path.getctime)
print("Resumed from", checkpoint_file)
checkpoint = torch.load(checkpoint_file, map_location='cpu')
model.load_state_dict(checkpoint['model'])
train_loss, valid_loss = checkpoint["train_loss"], checkpoint["valid_loss"]
valid_metric = checkpoint['valid_metric']
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['lr_scheduler'])
epoch_start = checkpoint["epoch"] + 1
early_stopping.load_state_dict(checkpoint["stopping_params"])
except ValueError:
print(f"Fold {fold_n + 1} checkpoint not found. Starting from scratch")
epoch_start = 0
else:
epoch_start = 0
if torch.cuda.is_available():
model, optimizer = apex.amp.initialize(model, optimizer, opt_level="O1")
if args.parallel:
model = torch.nn.DataParallel(model)
print(f'Fold {fold_n + 1} started at {time.ctime()}')
for epoch in range(epoch_start, args.epochs):
model.train()
running_loss = 0
for i, (image, target) in enumerate(train_loader):
# # https://www.tensorflow.org/versions/r1.14/api_docs/python/tf/losses/sigmoid_cross_entropy
# if args.label_smoothing > 0:
# target = target * (1 - args.label_smoothing) + 0.5 * args.label_smoothing
# # target[:, 1:] = target[:, 1:] * (1 - args.label_smoothing) + 0.5 * args.label_smoothing
image, target = image.to(device), target.to(device)
optimizer.zero_grad()
y_hat = model(image)
loss = weighted_multi_label_logloss(criterion, y_hat, target, weights_norm)
if torch.cuda.is_available():
with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 0 or i+1 == len(train_loader):
print(f"Fold {fold_n+1} [{epoch:2d}, {i+1:5d}/{len(train_loader):5d}] Loss: {loss.item():.4f}")
valid_predictions, valid_target = test_time_augmentation(model, valid_loaders, device, "Valid")
score = roc_auc_score(valid_target.numpy(), valid_predictions.numpy())
validation_loss = weighted_multi_label_logloss(criterion, valid_predictions,
valid_target, weights_norm).item()
# validation_metric = competition_metric(valid_target, valid_predictions)
validation_metric = validation_loss
running_loss = running_loss / len(train_loader)
print(Fore.GREEN + f'Train Loss: {running_loss:.5f}, '
f'Validation Loss: {validation_loss:.5f}, '
f'Validation Metric: {validation_metric:.5f}, '
f'ROC AUC: {score:.5f}')
train_loss.append(running_loss)
valid_loss.append(validation_loss)
valid_metric.append(validation_metric)
# scheduler.step(validation_metric)
scheduler.step()
early_stopping(validation_metric, model, epoch=epoch, args=args, optimizer=optimizer.state_dict(),
train_loss=train_loss, valid_loss=valid_loss, valid_metric=valid_metric,
lr_scheduler=scheduler.state_dict())
if early_stopping.early_stop:
print("Early stopping")
break
print(Fore.GREEN + f"Best score: {-early_stopping.best_score:.5f}")
fold_data[f"train_loss_{fold_n + 1}"] = pd.Series(train_loss)
fold_data[f"valid_loss_{fold_n + 1}"] = pd.Series(valid_loss)
fold_data[f"valid_metric_{fold_n + 1}"] = pd.Series(valid_metric)
# Predict on test and OOF sets using the best checkpoint on a fresh model
# For some reason loading weights to an existing model was doing weird things with the submission ¯\_(ツ)_/¯
print("Loading checkpoint and creating submission & OOFs")
model = model_builder(args.architecture)
model = model.to(device)
checkpoint = torch.load(checkpoint_file, map_location='cpu')
model.load_state_dict(checkpoint['model'])
if args.parallel:
model = torch.nn.DataParallel(model)
# scores.append(score)
valid_predictions, valid_target = test_time_augmentation(model, valid_loaders, device, "Valid")
oof_df.loc[valid_image_ids, categories] = torch.sigmoid(valid_predictions).numpy()
plot_multiclass_roc_curve(valid_target, torch.sigmoid(valid_predictions).numpy(),
os.path.join(args.output_dir, f"roc_fold_{fold_n + 1}.png"),
-early_stopping.best_score)
test_predictions, _ = test_time_augmentation(model, test_loaders, device, "Test ")
prediction += test_predictions.detach().numpy()
# Save results
test_df = pd.DataFrame(torch.sigmoid(torch.tensor(prediction / (fold_n + 1))).numpy(), columns=categories)
test_df["ImageID"] = test_loaders[0].dataset.image_ids
test_df.to_csv(os.path.join(args.output_dir, "test_predictions.csv"), index=False)
oof = wide_to_long(oof_df.reset_index(), id_var="SOPInstanceUID")
sub = reindex_submission(wide_to_long(test_df), args.stage)
oof.to_csv(os.path.join(args.output_dir, f"{args.output_dir[-15:]}_oof_predictions.csv.gz"),
index=False, compression="gzip")
sub.to_csv(os.path.join(args.output_dir, f"{args.output_dir[-15:]}_submission.csv.gz"),
index=False, compression="gzip")
# break # HACK: Single fold
metric_values = [v.iloc[-1] for k, v in fold_data.items() if 'valid_metric' in k]
for i, (ll, roc) in enumerate(zip(metric_values, scores)):
print(Fore.GREEN + f'Fold {i+1} LogLoss: {ll:.4f}, ROC AUC: {roc:.4f}')
print(Fore.CYAN + f'CV mean ROC AUC score: {np.mean(scores):.4f}, std: {np.std(scores):.4f}')
print(Fore.CYAN + f'CV mean Log Loss: {np.mean(metric_values):.4f}, std: {np.std(metric_values):.4f}')
# Make a plot
fold_data = pd.DataFrame(fold_data)
fig, ax = plt.subplots()
plt.title(f"CV Mean score: {np.mean(metric_values):.4f} +/- {np.std(metric_values):.4f}")
valid_curves = fold_data.loc[:, fold_data.columns.str.startswith('valid_loss')]
train_curves = fold_data.loc[:, fold_data.columns.str.startswith('train_loss')]
valid_curves.plot(ax=ax, colormap='Blues_r')
train_curves.plot(ax=ax, colormap='Reds_r')
# ax.set_ylim([np.min(train_curves.values), np.max(valid_curves.values)])
ax.tick_params(labelleft=True, labelright=True, left=True, right=True)
plt.savefig(os.path.join(args.output_dir, f"phase0.png"))
print("Done in", (time.time() - start) // 60, "minutes")
if __name__ == '__main__':
import argparse
import yaml
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--config', action='store', dest='config', help='Configuration scheme')
args = parser.parse_args()
# Lookup the config from the YAML file
with open("config.yml", 'r') as ymlfile:
cfg = yaml.load(ymlfile, Loader=yaml.FullLoader)
if args.config is None:
settings = cfg['default_run']
# settings = cfg['default_run_1080']
else:
print("Using", args.config, "configuration")
settings = cfg[args.config]
args.epochs = settings["epochs"]
args.batch = settings["batch"]
args.lr = settings["lr"]
args.img_size = settings["img_size"]
args.checkpoint = settings["checkpoint"]
args.parallel = settings["parallel"]
args.img_folder = settings["img_folder"]
args.architecture = settings["architecture"]
args.stage = settings["stage"]
args.cv_scheme = settings["cv_scheme"]
if args.checkpoint is None:
args.checkpoint = datetime.now().strftime("%Y%m%d-%H%M%S")
args.output_dir = os.path.join(OUTPUT_DIR, "phase0", args.checkpoint)
args.checkpoint_path = None
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
else:
args.checkpoint_path = os.path.join(OUTPUT_DIR, "phase0", args.checkpoint)
args.output_dir = os.path.join(OUTPUT_DIR, "phase0", datetime.now().strftime("%Y%m%d-%H%M%S"))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
print("Saving output to", args.output_dir)
copyfile("config.yml", os.path.join(args.output_dir, "config.yml"))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_model(args, device)
Datasets
Models
