import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.optim.optimizer
from torch.utils.data import DataLoader
from data.hyperkvasir import KvasirSegmentationDataset, KvasirMNVset
from evaluation.metrics import iou
from losses.consistency_losses import *
from perturbation.model import ModelOfNaturalVariation
from training.vanilla_trainer import VanillaTrainer
from utils import logging
from training.consistency_trainers import ConsistencyTrainer
from models.segmentation_models import InductiveNet
from models.ensembles import TrainedEnsemble
from data.hyperkvasir import KvasirSegmentationDataset, KvasirMNVset
from evaluation.metrics import iou
from losses.consistency_losses import *
from perturbation.model import ModelOfNaturalVariation
from training.vanilla_trainer import VanillaTrainer
from utils import logging
from data.etis import EtisDataset
class InductiveNetConsistencyTrainer:
def __init__(self, id, config):
"""
:param model: String describing the model type. Can be DeepLab, TriUnet, ... TODO
:param config: Contains hyperparameters : lr, epochs, batch_size, T_0, T_mult
"""
self.config = config
self.device = config["device"]
self.lr = config["lr"]
self.batch_size = config["batch_size"]
self.epochs = config["epochs"]
self.id = id
self.model_str = "InductiveNet"
self.mnv = ModelOfNaturalVariation(T0=1).to(self.device)
self.nakedcloss = NakedConsistencyLoss()
self.closs = ConsistencyLoss()
self.model = InductiveNet().to(self.device)
self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr)
self.jaccard = vanilla_losses.JaccardLoss()
self.mse = nn.MSELoss()
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=50, T_mult=2)
self.train_set = KvasirSegmentationDataset("Datasets/HyperKvasir", split="train", augment=False)
self.val_set = KvasirSegmentationDataset("Datasets/HyperKvasir", split="val")
self.test_set = KvasirSegmentationDataset("Datasets/HyperKvasir", split="test")
self.train_loader = DataLoader(self.train_set, batch_size=self.batch_size, shuffle=True)
self.val_loader = DataLoader(self.val_set)
self.test_loader = DataLoader(self.test_set)
def train_epoch(self):
self.model.train()
losses = []
for x, y, fname in self.train_loader:
image = x.to("cuda")
mask = y.to("cuda")
aug_img, aug_mask = self.mnv(image, mask)
self.optimizer.zero_grad()
aug_output, _ = self.model(aug_img)
output, reconstruction = self.model(image)
mean_iou = torch.mean(iou(output, mask))
loss = 0.5 * (self.closs(aug_mask, mask, aug_output, output, mean_iou) + self.mse(
image, reconstruction))
loss.backward()
self.optimizer.step()
losses.append(np.abs(loss.item()))
return np.mean(losses)
def train(self):
best_val_loss = 10
print("Starting Segmentation training")
best_consistency = 0
for i in range(self.epochs):
training_loss = np.abs(self.train_epoch())
val_loss, ious, closs = self.validate(epoch=i, plot=False)
gen_ious = self.validate_generalizability(epoch=i, plot=False)
mean_iou = float(torch.mean(ious))
gen_iou = float(torch.mean(gen_ious))
consistency = 1 - np.mean(closs)
test_iou = np.mean(self.test().numpy())
self.config["lr"] = [group['lr'] for group in self.optimizer.param_groups]
logging.log_full(epoch=i, id=self.id, config=self.config, result_dict=
{"train_loss": training_loss, "val_loss": val_loss,
"iid_val_iou": mean_iou, "iid_test_iou": test_iou, "ood_iou": gen_iou,
"consistency": consistency}, type="consistency")
self.scheduler.step(i)
print(
f"Epoch {i} of {self.epochs} \t"
f" lr={[group['lr'] for group in self.optimizer.param_groups]} \t"
f" loss={training_loss} \t"
f" val_loss={val_loss} \t"
f" ood_iou={gen_iou}\t"
f" val_iou={mean_iou} \t"
f" gen_prop={gen_iou / mean_iou}"
)
if val_loss < best_val_loss:
best_val_loss = val_loss
np.save(
f"experiments/Data/Augmented-Pipelines/{self.model_str}/{self.id}",
test_iou)
print(f"Saving new best model. IID test-set mean iou: {test_iou}")
torch.save(self.model.state_dict(),
f"Predictors/Augmented/{self.model_str}/{self.id}")
print("saved in: ", f"Predictors/Augmented/{self.model_str}/{self.id}")
if consistency > best_consistency:
best_consistency = consistency
torch.save(self.model.state_dict(),
f"Predictors/Augmented/{self.model_str}/maximum_consistency{self.id}")
torch.save(self.model.state_dict(),
f"Predictors/Augmented/{self.model_str}/{self.id}_last_epoch")
def test(self):
self.model.eval()
ious = torch.empty((0,))
with torch.no_grad():
for x, y, fname in self.test_loader:
image = x.to("cuda")
mask = y.to("cuda")
output, _ = self.model(image)
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
ious = torch.cat((ious, batch_ious.flatten()))
return ious
def validate(self, epoch, plot=False):
self.model.eval()
losses = []
closses = []
ious = torch.empty((0,))
with torch.no_grad():
for x, y, fname in self.val_loader:
image = x.to("cuda")
mask = y.to("cuda")
aug_img, aug_mask = self.mnv(image, mask)
output, reconstruction = self.model(image)
aug_output, _ = self.model(aug_img) # todo consider train on augmented vs non-augmented?
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
loss = 0.5 * (self.closs(aug_mask, mask, aug_output, output, torch.mean(batch_ious)) + self.mse(
image, reconstruction))
losses.append(np.abs(loss.item()))
closses.append(self.nakedcloss(aug_mask, mask, aug_output, output).item())
ious = torch.cat((ious, batch_ious.cpu().flatten()))
if plot:
plt.imshow(output[0, 0].cpu().numpy(), alpha=0.5)
plt.imshow(reconstruction[0].permute(1, 2, 0).cpu().numpy())
# plt.imshow((output[0, 0].cpu().numpy() > 0.5).astype(int), alpha=0.5)
# plt.imshow(y[0, 0].cpu().numpy().astype(int), alpha=0.5)
# plt.title("IoU {} at epoch {}".format(iou(output[0, 0], mask[0, 0]), epoch))
plt.show()
plot = False # plot one example per epoch
avg_val_loss = np.mean(losses)
avg_closs = np.mean(closses)
return avg_val_loss, ious, closses
def validate_generalizability(self, epoch, plot=False):
self.model.eval()
ious = torch.empty((0,))
with torch.no_grad():
for x, y, index in DataLoader(EtisDataset("Datasets/ETIS-LaribPolypDB")):
image = x.to("cuda")
mask = y.to("cuda")
output, _ = self.model(image)
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
ious = torch.cat((ious, batch_ious.flatten()))
if plot:
plt.imshow(image[0].permute(1, 2, 0).cpu().numpy())
plt.imshow((output[0, 0].cpu().numpy() > 0.5).astype(int), alpha=0.5)
plt.title("IoU {} at epoch {}".format(iou(output[0, 0], mask[0, 0]), epoch))
plt.show()
plot = False # plot one example per epoch (hacky, but works)
return ious
class InductiveNetVanillaTrainer(InductiveNetConsistencyTrainer):
def __init__(self, id, config):
super(InductiveNetVanillaTrainer, self).__init__(id, config)
def train(self):
best_val_loss = 10
print("Starting Segmentation training")
best_consistency = 0
for i in range(self.epochs):
training_loss = np.abs(self.train_epoch())
val_loss, ious, closs = self.validate(epoch=i, plot=False)
gen_ious = self.validate_generalizability(epoch=i, plot=False)
mean_iou = float(torch.mean(ious))
gen_iou = float(torch.mean(gen_ious))
consistency = 1 - np.mean(closs)
test_iou = np.mean(self.test().numpy())
self.config["lr"] = [group['lr'] for group in self.optimizer.param_groups]
logging.log_full(epoch=i, id=self.id, config=self.config, result_dict=
{"train_loss": training_loss, "val_loss": val_loss,
"iid_val_iou": mean_iou, "iid_test_iou": test_iou, "ood_iou": gen_iou,
"consistency": consistency}, type="consistency")
self.scheduler.step(i)
print(
f"Epoch {i} of {self.epochs} \t"
f" lr={[group['lr'] for group in self.optimizer.param_groups]} \t"
f" loss={training_loss} \t"
f" val_loss={val_loss} \t"
f" ood_iou={gen_iou}\t"
f" val_iou={mean_iou} \t"
f" gen_prop={gen_iou / mean_iou}"
)
if val_loss < best_val_loss:
best_val_loss = val_loss
np.save(
f"experiments/Data/Normal-Pipelines/{self.model_str}/{self.id}",
test_iou)
print(f"Saving new best model. IID test-set mean iou: {test_iou}")
torch.save(self.model.state_dict(),
f"Predictors/Vanilla/{self.model_str}/{self.id}")
print("saved in: ", f"Predictors/Vanilla/{self.model_str}/{self.id}")
if consistency > best_consistency:
best_consistency = consistency
torch.save(self.model.state_dict(),
f"Predictors/Vanilla/{self.model_str}/maximum_consistency{self.id}")
torch.save(self.model.state_dict(),
f"Predictors/Vanilla/{self.model_str}/{self.id}_last_epoch")
def train_epoch(self):
self.model.train()
losses = []
for x, y, fname in self.train_loader:
image = x.to("cuda")
mask = y.to("cuda")
self.optimizer.zero_grad()
output, reconstruction = self.model(image)
mean_iou = torch.mean(iou(output, mask))
loss = 0.5 * (self.jaccard(output, mask) + self.mse(
image, reconstruction))
loss.backward()
self.optimizer.step()
losses.append(np.abs(loss.item()))
return np.mean(losses)
def test(self):
self.model.eval()
ious = torch.empty((0,))
with torch.no_grad():
for x, y, fname in self.test_loader:
image = x.to("cuda")
mask = y.to("cuda")
output = self.model(image)
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
ious = torch.cat((ious, batch_ious.flatten()))
return ious
def validate(self, epoch, plot=False):
self.model.eval()
losses = []
closses = []
ious = torch.empty((0,))
with torch.no_grad():
for x, y, fname in self.val_loader:
image = x.to("cuda")
mask = y.to("cuda")
aug_img, aug_mask = self.mnv(image, mask)
output, reconstruction = self.model(image)
aug_output, _ = self.model(aug_img) # todo consider train on augmented vs non-augmented?
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
loss = 0.5 * (self.jaccard(output, mask) + self.mse(
image, reconstruction))
losses.append(np.abs(loss.item()))
closses.append(self.nakedcloss(aug_mask, mask, aug_output, output).item())
ious = torch.cat((ious, batch_ious.cpu().flatten()))
if plot:
plt.imshow(output[0, 0].cpu().numpy(), alpha=0.5)
plt.imshow(reconstruction[0].permute(1, 2, 0).cpu().numpy())
# plt.imshow((output[0, 0].cpu().numpy() > 0.5).astype(int), alpha=0.5)
# plt.imshow(y[0, 0].cpu().numpy().astype(int), alpha=0.5)
# plt.title("IoU {} at epoch {}".format(iou(output[0, 0], mask[0, 0]), epoch))
plt.show()
plot = False # plot one example per epoch
avg_val_loss = np.mean(losses)
avg_closs = np.mean(closses)
return avg_val_loss, ious, avg_closs
def validate_generalizability(self, epoch, plot=False):
self.model.eval()
ious = torch.empty((0,))
with torch.no_grad():
for x, y, index in DataLoader(EtisDataset("Datasets/ETIS-LaribPolypDB")):
image = x.to("cuda")
mask = y.to("cuda")
output, _ = self.model(image)
batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
ious = torch.cat((ious, batch_ious.flatten()))
if plot:
plt.imshow(image[0].permute(1, 2, 0).cpu().numpy())
plt.imshow((output[0, 0].cpu().numpy() > 0.5).astype(int), alpha=0.5)
plt.title("IoU {} at epoch {}".format(iou(output[0, 0], mask[0, 0]), epoch))
plt.show()
plot = False # plot one example per epoch (hacky, but works)
return ious
class InductiveNetAugmentationTrainer(InductiveNetConsistencyTrainer):
"""
Uses vanilla data augmentation with p=0.5 instead of a a custom loss
"""
def __init__(self, id, config):
super(InductiveNetAugmentationTrainer, self).__init__(id, config)
self.jaccard = vanilla_losses.JaccardLoss()
self.mse = vanilla_losses.MSELoss()
self.prob = 0
self.dataset = KvasirMNVset("Datasets/HyperKvasir", "train", inpaint=config["use_inpainter"])
self.train_loader = DataLoader(self.dataset, batch_size=config["batch_size"], shuffle=True)
def get_iou_weights(self, image, mask):
self.model.eval()
with torch.no_grad():
output, _ = self.model(image)
return torch.mean(iou(output, mask))
def get_consistency(self, image, mask, augmented, augmask):
self.model.eval()
with torch.no_grad():
output, _ = self.model(image)
self.model.train()
return torch.mean(self.nakedcloss(output, mask, augmented, augmask))
def train_epoch(self):
self.model.train()
losses = []
for x, y, fname, flag in self.train_loader:
image = x.to("cuda")
mask = y.to("cuda")
self.optimizer.zero_grad()
output, reconstruction = self.model(image)
mean_iou = torch.mean(iou(output, mask))
loss = 0.5 * (self.jaccard(output, mask) + self.mse(
image, reconstruction))
loss.backward()
self.optimizer.step()
losses.append(np.abs(loss.item()))
return np.mean(losses)
class InductiveNetEnsembleTrainer(InductiveNetConsistencyTrainer):
def __init__(self, id, config):
super(InductiveNetEnsembleTrainer, self).__init__(id, config)
self.model = TrainedEnsemble("Singular")
self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr)
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=50, T_mult=2)
Datasets
Models
