import warnings
warnings.simplefilter("ignore")
import math
from typing import Any, Union
import torch
import torch.utils.checkpoint as torch_ckp
import torchvision
import torchmetrics
import torchmetrics.classification
import lightning as L
from enum import Enum
import timm
import van
import numpy as np
from softadapt import LossWeightedSoftAdapt, NormalizedSoftAdapt
from noah import NOAH
import piq
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import LabelEncoder
from tqdm import tqdm
import pandas as pd
from pathlib import Path
from I3D_RestNet50 import I3DResNet50
def js_divergence_sigmoid(p_logits, q_logits):
p_probs = torch.sigmoid(p_logits)
q_probs = torch.sigmoid(q_logits)
m = 0.5 * (p_probs + q_probs)
bce_p_m = torch.nn.functional.binary_cross_entropy(m, p_probs, reduction='none')
bce_q_m = torch.nn.functional.binary_cross_entropy(m, q_probs, reduction='none')
js_div = 0.5 * (bce_p_m + bce_q_m)
return js_div
def getNVL(obj:dict, key, default):
return default if obj.get(key) is None else obj.get(key)
class Decoder (torch.nn.Module):
def __init__(self, input_channels=2048, output_channels=3, num_blocks=4, hidden_channels=[1024, 512, 256, 64], *args, **kwargs):
super().__init__(*args, **kwargs)
assert len(hidden_channels) == num_blocks, "Number of hidden channels must match the number of blocks."
self.Activation = torch.nn.ReLU#torch.nn.GELU
self.input_channels = input_channels
blocks = []
inCh = input_channels
blocks.append(torch.nn.Conv2d(inCh, inCh, kernel_size=3, stride=1, padding=1, bias=False))
blocks.append(torch.nn.BatchNorm2d(inCh))
blocks.append(self.Activation())
for i, outCh in enumerate(hidden_channels):
blocks.append(torch.nn.ConvTranspose2d(inCh, outCh, kernel_size=4, stride=2, padding=1, bias=False))
blocks.append(torch.nn.BatchNorm2d(outCh))
blocks.append(self.Activation())
blocks.append(torch.nn.Conv2d(outCh, outCh, kernel_size=3, stride=1, padding=1, bias=False))
blocks.append(torch.nn.BatchNorm2d(outCh))
blocks.append(self.Activation())
#blocks.append(torch.nn.ConvTranspose2d(inCh, outCh, kernel_size=3, stride=2, padding=1, output_padding=1))
#blocks.append(torch.nn.BatchNorm2d(outCh))
#blocks.append(self.Activation())
inCh = outCh
blocks.append(torch.nn.ConvTranspose2d(inCh, output_channels, kernel_size=4, stride=2, padding=1, bias=False))
blocks.append(torch.nn.BatchNorm2d(output_channels))
blocks.append(self.Activation())
blocks.append(torch.nn.Conv2d(output_channels, output_channels, kernel_size=3, stride=1, padding=1, bias=False))
blocks.append(torch.nn.BatchNorm2d(output_channels))
blocks.append(self.Activation())
#blocks.append(torch.nn.ConvTranspose2d(inCh, output_channels, kernel_size=3, stride=2, padding=1, output_padding=1))
self.decoder = torch.nn.Sequential(*blocks)
def forward(self, x):
x = self.decoder(x)
return x
class DynamicDecoder(torch.nn.Module):
def __init__(self, input_channels=2048, output_channels=3, num_blocks=4, hidden_channels=[1024, 512, 256, 64]):
super(DynamicDecoder, self).__init__()
# Ensure the number of hidden channels matches the number of blocks
assert len(hidden_channels) == num_blocks, "Number of hidden channels must match the number of blocks."
layers = []
in_channels = input_channels
# Loop to create the decoder blocks
for out_channels in hidden_channels:
layers.append(torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1))
layers.append(torch.nn.BatchNorm2d(out_channels))
layers.append(torch.nn.ReLU(inplace=True))
in_channels = out_channels
# Final layer to get the output image
layers.append(torch.nn.ConvTranspose2d(in_channels, output_channels, kernel_size=3, stride=2, padding=1, output_padding=1))
#layers.append(torch.nn.Sigmoid()) # To get pixel values between 0 and 1
# Combine all layers into a Sequential module
self.decoder = torch.nn.Sequential(*layers)
def forward(self, x):
return self.decoder(x)
class DecoderUnet(torch.nn.Module):
def __init__(self, input_channels=2048, output_channels=3):
super().__init__()
self.dropout = torch.nn.Dropout2d(0.4)
self.upConv_0 = torch.nn.ConvTranspose2d(input_channels, 512, kernel_size=2, stride=2)
self.decoder_0 = self._conv_block(1024, 512)
self.upConv_1 = torch.nn.ConvTranspose2d(512, 320, kernel_size=2, stride=2)
self.decoder_1 = self._conv_block(640, 320)
self.upConv_2 = torch.nn.ConvTranspose2d(320, 128, kernel_size=2, stride=2)
self.decoder_2 = self._conv_block(256, 128)
self.upConv_3 = torch.nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
self.decoder_3 = self._conv_block(128, 64)
self.upConv_4 = torch.nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
self.decoder_4 = self._conv_block(32, 16)
self.enc3_upsample = torch.nn.ConvTranspose2d(512, 512, kernel_size=2, stride=2)
self.enc2_upsample = torch.nn.ConvTranspose2d(320, 320, kernel_size=2, stride=2)
self.enc1_upsample = torch.nn.ConvTranspose2d(128, 128, kernel_size=2, stride=2)
self.enc0_upsample = torch.nn.ConvTranspose2d(64, 64, kernel_size=2, stride=2)
self.last_conv = torch.nn.Conv2d(16, output_channels, kernel_size=1)
def _conv_block(self, in_ch, out_ch):
return torch.nn.Sequential(
torch.nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
torch.nn.SiLU(inplace=True),
torch.nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
torch.nn.SiLU(inplace=True),
)
def forward(self, x):
encoder_l0, encoder_l1, encoder_l2, encoder_l3, btlneck = x
decoder_l3 = self.upConv_0(btlneck)
encoder_l3 = self.enc3_upsample(encoder_l3)
decoder_l3 = torch.cat((decoder_l3, encoder_l3), dim=1)
decoder_l3 = self.decoder_0(decoder_l3)
decoder_l3 = self.dropout(decoder_l3)
decoder_l2 = self.upConv_1(decoder_l3)
encoder_l2 = self.enc2_upsample(encoder_l2)
decoder_l2 = torch.cat((decoder_l2, encoder_l2), dim=1)
decoder_l2 = self.decoder_1(decoder_l2)
decoder_l2 = self.dropout(decoder_l2)
decoder_l1 = self.upConv_2(decoder_l2)
encoder_l1 = self.enc1_upsample(encoder_l1)
decoder_l1 = torch.cat((decoder_l1, encoder_l1), dim=1)
decoder_l1 = self.decoder_2(decoder_l1)
decoder_l1 = self.dropout(decoder_l1)
decoder_l0 = self.upConv_3(decoder_l1)
encoder_l0 = self.enc0_upsample(encoder_l0)
decoder_l0 = torch.cat((decoder_l0, encoder_l0), dim=1)
decoder_l0 = self.decoder_3(decoder_l0)
decoder_last = self.upConv_4(decoder_l0)
decoder_last = self.decoder_4(decoder_last)
return self.last_conv(decoder_last)
class ModelsList(Enum):
RestNet50 = 1
DenseNet169 = 2
Efficientnet_b0 = 3
MobileNetV2 = 4
Inception3 = 5
ResNeXt_50_32x4d = 6
RestNet50_Droput = 7
class TypeLossFunction(Enum):
CrossEntropy = 0
BCEWithLogits = 1
HingeLoss = 2
FocalLoss = 3
ContrastiveLoss = 4
class FeatureAlignmentLoss(torch.nn.Module):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def forward(self, F_student, F_teacher):
fS = torch.nn.functional.normalize(F_student, p=2, dim=-1)
fT = torch.nn.functional.normalize(F_teacher, p=2, dim=-1)
cos_Sim = torch.sum(fS * fT, dim=-1)
loss = 1 - cos_Sim
return loss.mean()
class CosineAnnealingLR(torch.optim.lr_scheduler._LRScheduler):
def __init__(
self,
optimizer: torch.optim.Optimizer,
warmup_epochs: int,
max_epochs: int,
warmup_start_lr: float = 0.00001,
eta_min: float = 0.00001,
last_epoch: int = -1,
):
"""
Args:
optimizer (torch.optim.Optimizer):
最適化手法インスタンス
warmup_epochs (int):
linear warmupを行うepoch数
max_epochs (int):
cosine曲線の終了に用いる 学習のepoch数
warmup_start_lr (float):
linear warmup 0 epoch目の学習率
eta_min (float):
cosine曲線の下限
last_epoch (int):
cosine曲線の位相オフセット
学習率をmax_epochsに至るまでコサイン曲線に沿ってスケジュールする
epoch 0からwarmup_epochsまでの学習曲線は線形warmupがかかる
https://pytorch-lightning-bolts.readthedocs.io/en/stable/schedulers/warmup_cosine_annealing.html
"""
self.warmup_epochs = warmup_epochs
self.max_epochs = max_epochs
self.warmup_start_lr = warmup_start_lr
self.eta_min = eta_min
super().__init__(optimizer, last_epoch, verbose=False)
return None
def get_lr(self):
if self.last_epoch == 0:
return [self.warmup_start_lr] * len(self.base_lrs)
if self.last_epoch < self.warmup_epochs:
return [
group["lr"] + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
]
if self.last_epoch == self.warmup_epochs:
return self.base_lrs
if (self.last_epoch - 1 - self.max_epochs) % (2 * (self.max_epochs - self.warmup_epochs)) == 0:
return [
group["lr"] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs))) / 2
for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
]
return [
(1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)))
/ (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs - 1) / (self.max_epochs - self.warmup_epochs)))
* (group["lr"] - self.eta_min)
+ self.eta_min
for group in self.optimizer.param_groups
]
class ContrastiveLoss(torch.nn.Module):
def __init__(self, margin=1.0, distance:int = 0):
super(ContrastiveLoss, self).__init__()
self.margin = margin
self.TypeDistance = distance
def forward(self, output1, output2, label):
# Calcula la distancia euclidiana entre las dos salidas
Distance = torch.nn.functional.pairwise_distance(output1, output2)
# Calcula la pérdida contrastiva
loss_contrastive = torch.mean(
(1 - label) * torch.pow(Distance, 2) +
(label) * torch.pow(torch.clamp(self.margin - Distance, min=0.0), 2)
)
return loss_contrastive
class GradCAM:
def __init__(self, model, target_layer):
self.model = model
self.target_layer = target_layer
self.gradients = None
self.activations = None
self.hook()
def hook(self):
def forward_hook(module, Input, Output):
self.activations = Output
def backward_hook(module, grad_in, grad_out):
self.gradients = grad_out[0]
self.target_layer.register_forward_hook(forward_hook)
self.target_layer.register_backward_hook(backward_hook)
def generate_cam(self, data:torch.Tensor, target_class):
self.model.zero_grad()
data.requires_grad= True
output = self.model(data).flatten()
loss = torch.nn.functional.binary_cross_entropy_with_logits(output, target_class)
loss.backward()
pooled_Grad = torch.mean(self.gradients, dim=[0, 2, 3])
for i in range (pooled_Grad.size(0)):
self.activations[:, i, :, :] *= pooled_Grad[i]
cam = torch.mean(self.activations, dim=1).squeeze()
cam = np.maximum(cam.detach().numpy(), 0)
cam = (cam - cam.min()) / (cam.max() - cam.min())
return cam
class UNet_RN18(torch.nn.Module):
def _conv_block(self, in_ch, out_ch):
return torch.nn.Sequential(
torch.nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
torch.nn.SiLU(inplace=True),
torch.nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
torch.nn.SiLU(inplace=True),
)
def _hook_fn(self, module, input, output):
self.feature_maps.append(output)
def _register_encoder_hooks(self):
for layer in self.list_blocks:
self.hooks.append(layer.register_forward_hook(self._hook_fn))
def __init__(self, in_channels:int = 3, out_channels:int = 1, *args, **kwargs):
super().__init__(*args, **kwargs)
self.hooks = []
self.feature_maps = []
self.encoder_base = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
self.encoder_base.fc = torch.nn.Identity()
#for parms in self.encoder_base.parameters():
# parms.requires_grad = False
self.list_blocks = [
self.encoder_base.conv1,
self.encoder_base.layer1,
self.encoder_base.layer2,
self.encoder_base.layer3,
self.encoder_base.layer4
]
self._register_encoder_hooks()
self.dropout = torch.nn.Dropout2d(0.4)
self.upConv_0 = torch.nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
self.upConv_1 = torch.nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
self.upConv_2 = torch.nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
self.upConv_3 = torch.nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
self.upConv_extra = torch.nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
self.decoder_0 = self._conv_block(512, 256) #0
self.decoder_1 = self._conv_block(256, 128) #1
self.decoder_2 = self._conv_block(128, 64) #2
self.decoder_3 = self._conv_block(32 + 64, 32) #3
self.decoder_extra = self._conv_block(16, 16)
self.last_conv = torch.nn.Conv2d(16, out_channels, kernel_size=1)
def forward(self, x):
self.feature_maps = []
_ = self.encoder_base(x) # forward to encoder and call hooks
encoder_l0, encoder_l1, encoder_l2, encoder_l3, btlneck = self.feature_maps
decoder_l3 = self.upConv_0(btlneck)
decoder_l3 = torch.cat((decoder_l3, encoder_l3), dim=1)
decoder_l3 = self.decoder_0(decoder_l3)
decoder_l3 = self.dropout(decoder_l3)
decoder_l2 = self.upConv_1(decoder_l3)
decoder_l2 = torch.cat((decoder_l2, encoder_l2), dim=1)
decoder_l2 = self.decoder_1(decoder_l2)
decoder_l2 = self.dropout(decoder_l2)
decoder_l1 = self.upConv_2(decoder_l2)
decoder_l1 = torch.cat((decoder_l1, encoder_l1), dim=1)
decoder_l1 = self.decoder_2(decoder_l1)
decoder_l1 = self.dropout(decoder_l1)
decoder_l0 = self.upConv_3(decoder_l1)
decoder_l0 = torch.cat((decoder_l0, encoder_l0), dim=1)
decoder_l0 = self.decoder_3(decoder_l0)
decoder_last = self.upConv_extra(decoder_l0)
decoder_last = self.decoder_extra(decoder_last)
return self.last_conv(decoder_last)
class RARP_NVB_ROI_Mask_Unet(L.LightningModule):
def __init__(self,*args, **kwargs):
super().__init__(*args, **kwargs)
self.model = UNet_RN18(in_channels=3, out_channels=1)
self.lr = 1E-4
self.Lambda_L1 = None
self.lossFN = torch.nn.BCEWithLogitsLoss()
self.train_IoU = torchmetrics.classification.BinaryJaccardIndex()
self.val_IoU = torchmetrics.classification.BinaryJaccardIndex()
def forward(self, data):
data = data.float()
pred = self.model(data)
return pred
def _shared_step(self, batch, val_step:bool = True):
img, mask = batch
mask = mask.float()
mask = mask.unsqueeze(1)
prediction = self(img)
loss = self.lossFN(prediction, mask)
predicted_labels = torch.sigmoid(prediction)
if not val_step:
if self.Lambda_L1 is not None:
loss_l1 = 0
for name, params in self.model.named_parameters():
if "decoder" in name or "upConv" in name:
loss_l1 += torch.norm(params, p=1)
loss += self.Lambda_L1 * loss_l1
return loss, mask, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch, False)
self.train_IoU.update(predicted_labels, true_labels)
self.log("train_loss", loss, on_epoch=True)
self.log("train_acc_IoU", self.train_IoU, on_epoch=True, on_step=False)
return loss
def on_after_backward(self):
total_norm = 0.0
for p in self.parameters():
if p.grad is not None:
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
self.log("grad_norm", total_norm)
if total_norm < 1e-8:
self.log("grad_warning", "Vanishing gradient suspected!")
def on_train_epoch_start(self):
for parms in self.model.encoder_base.parameters():
parms.requires_grad = (self.current_epoch % 2 == 0)
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.val_IoU.update(predicted_labels, true_labels)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_acc_IoU", self.val_IoU, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
return [optimizer]
class RARP_NVB_ResNet50_CAM(L.LightningModule):
def __init__(self) -> None:
super().__init__()
self.model = torchvision.models.resnet50()
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.feature_map = torch.nn.Sequential(*list(self.model.children())[:-2])
def forward(self, data):
featureMap = self.feature_map(data)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(input=featureMap, output_size=(1, 1))
Cont_Net = torch.flatten(Cont_Net, 1)
pred = self.model.fc(Cont_Net)
return pred, featureMap
class RARP_NVB_VAN_CAM(L.LightningModule): #TODO
def __init__(self) -> None:
super().__init__()
self.model = van.van_b2(pretrained = True)
tempFC_ft = self.model.head.in_features
self.model.head = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.feature_map = torch.nn.Sequential(*list(self.model.children())[:-2])
def forward(self, data, label:torch.Tensor):
cams = GradCAM(self.model, self.model.block4[-1])
featureMap = cams.generate_cam(data, label)
pred = self.model(featureMap)
return pred, featureMap
class RARP_NVB_ResNet18_CAM(L.LightningModule):
def __init__(self) -> None:
super().__init__()
self.model = torchvision.models.resnet18()
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.feature_map = torch.nn.Sequential(*list(self.model.children())[:-2])
def forward(self, data):
featureMap = self.feature_map(data)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(input=featureMap, output_size=(1, 1))
Cont_Net = torch.flatten(Cont_Net)
pred = self.model.fc(Cont_Net)
return pred, featureMap
class RARP_NVB_MobileNetV2_CAM(L.LightningModule):
def __init__(self) -> None:
super().__init__()
self.model = torchvision.models.mobilenet_v2()
tempFC_ft = self.model.classifier[1].in_features
self.model.classifier[1] = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.feature_map = self.model.features
def forward(self, data):
featureMap = self.feature_map(data)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(featureMap, (1, 1))
Cont_Net = torch.flatten(Cont_Net, 1)
pred = self.model.classifier(Cont_Net)
return pred, featureMap
class RARP_NVB_EfficientNetV2_CAM(L.LightningModule):
def __init__(self) -> None:
super().__init__()
self.model = torchvision.models.efficientnet_v2_s(weights=torchvision.models.EfficientNet_V2_S_Weights.DEFAULT)
tempFC_ft = self.model.classifier[1].in_features
self.model.classifier[1] = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.feature_map = self.model.features
def forward(self, data):
featureMap = self.feature_map(data)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(featureMap, (1, 1))
Cont_Net = torch.flatten(Cont_Net, 1)
pred = self.model.classifier(Cont_Net)
return pred, featureMap
class RARP_NVB_Model_BCEWithLogitsLoss(L.LightningModule):
def __init__(self, x=None) -> None:
super().__init__()
self.model = None
self.lossFN = torch.nn.BCEWithLogitsLoss() #pos_weight=torch.tensor([2.73])
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1Score = torchmetrics.F1Score('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
def forward(self, data):
data = data.float()
return self.model(data)
def _shared_step(self, batch):
img, label = batch
label = label.float()
pred = self.forward(img).flatten()
loss = self.lossFN(pred, label)
predicted_labels = torch.sigmoid(pred)
return loss, label, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("train_loss", loss)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("val_loss", loss)
self.val_acc(predicted_labels, true_labels)
self.f1Score(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_f1", self.f1Score, on_epoch=True, on_step=False, prog_bar=True)
return loss
def test_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.test_acc(predicted_labels, true_labels)
self.f1ScoreTest(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
return [optimizer]
class RARP_NVB_FOCAL_loss(torch.nn.Module):
def __init__(self, alpha: float = 0.25, gamma: float = 2, reduction: str = "mean"):
super().__init__()
self.alpha = alpha
self.gamma = gamma
self.reduction = reduction
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return torchvision.ops.focal_loss.sigmoid_focal_loss(input, target, self.alpha, self.gamma, self.reduction)
class FocalLoss(torch.nn.Module):
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
super(FocalLoss, self).__init__()
self.alpha = alpha
self.gamma = gamma
self.reduction = reduction
def forward(self, inputs, targets):
# Convert targets to one-hot encoding
targets = torch.nn.functional.one_hot(targets, num_classes=inputs.size(1)).float()
# Compute softmax over the inputs
probs = torch.nn.functional.softmax(inputs, dim=1)
log_probs = torch.nn.functional.log_softmax(inputs, dim=1)
# Compute the focal loss components
focal_weight = (1 - probs) ** self.gamma
loss = -self.alpha * focal_weight * targets * log_probs
# Apply reduction method
if self.reduction == 'mean':
return loss.mean()
elif self.reduction == 'sum':
return loss.sum()
else:
return loss
#TODO
class RARP_NVB_MultiClassModel(L.LightningModule):
def __init__(self,
InitWeight = torch.tensor([1,1]),
schedulerLR:bool=False,
lr:float = 1e-4,
Model:torch.nn.Module = None,
Num_Classes:int = 2,
L1:float = 1.31E-04,
L2:float = 0
) -> None:
super().__init__()
self.model = Model
self.lossFN = FocalLoss() #torch.nn.CrossEntropyLoss()
self.InitWeight = InitWeight
self.scheduler = schedulerLR
self.lr = lr
self.Lambda_L1 = L1
self.Lambda_L2 = L2
self.num_classes = Num_Classes
self.train_acc = torchmetrics.Accuracy("multiclass", num_classes=Num_Classes)
self.val_acc = torchmetrics.Accuracy("multiclass", num_classes=Num_Classes)
self.test_acc = torchmetrics.Accuracy("multiclass", num_classes=Num_Classes)
self.f1ScoreTest = torchmetrics.F1Score("multiclass", num_classes=Num_Classes)
self.val_loop = False
def forward(self, data):
data = data.float()
pred = self.model(data)
return pred
def _shared_step(self, batch):
img, label = batch
prediction = self(img)
predicted_labels = torch.softmax(prediction, dim=1)
loss = self.lossFN(prediction, label)
if self.Lambda_L1 is not None and not self.val_loop:
loss_l1 = 0
for params in self.model.parameters():
loss_l1 += torch.sum(torch.abs(params))
loss += self.Lambda_L1 * loss_l1
return loss, label, predicted_labels
def training_step(self, batch, batch_idx):
self.val_loop = False
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
self.val_loop = True
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
self.val_loop = True
loss, true_labels, predicted_labels = self._shared_step(batch)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.Lambda_L2)
if self.scheduler:
#scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", patience=1, verbose=True, factor=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 6, eta_min=1e-8, verbose=True)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
#"monitor": "val_loss",
}
}
else:
return [optimizer]
class RARP_NVB_Model(L.LightningModule):
def __init__(self,
InitWeight = torch.tensor([1,1]),
typeLossFN:TypeLossFunction = TypeLossFunction.CrossEntropy,
schedulerLR:bool=True,
lr:float = 1e-4
) -> None:
super().__init__()
self.model = None
self.lossFN = torch.nn.CrossEntropyLoss(label_smoothing=0.5) if typeLossFN == TypeLossFunction.CrossEntropy else torch.nn.BCEWithLogitsLoss()
self.InitWeight = InitWeight
self.scheduler = schedulerLR
self.lr = lr
self.Lambda_L1 = None #1.31E-04 #
self.Lambda_L2 = 0
print (f"LR= {self.lr}, L1= {self.Lambda_L1}")
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
#self.f1Score = torchmetrics.F1Score('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
def forward(self, data):
data = data.float()
pred = self.model(data)
return pred
def _shared_step(self, batch):
img, label = batch
if self.InitWeight is not None:
self.lossFN.weight = self.InitWeight[label]
label = label.float()
prediction = self(img)[:,0] #.flatten()
loss = self.lossFN(prediction, label)
if self.Lambda_L1 is not None:
loss_l1 = 0
for params in self.parameters():
loss_l1 += torch.sum(torch.abs(params))
loss += self.Lambda_L1 * loss_l1
predicted_labels = torch.sigmoid(prediction)
return loss, label, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
#self.f1Score.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("hp_metric", self.val_acc, on_step=False, on_epoch=True)
#self.log("val_f1", self.f1Score, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.Lambda_L2)
if self.scheduler:
#scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", patience=1, verbose=True, factor=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 6, eta_min=1e-8, verbose=True)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
#"monitor": "val_loss",
}
}
else:
return optimizer
class RARP_Ensemble(L.LightningModule):
def __init__(self, ListModels,
InitWeight = torch.tensor([1,1]),
typeLossFN:TypeLossFunction = TypeLossFunction.CrossEntropy,
schedulerLR:bool=False,
lr:float = 1e-4) -> None:
super().__init__()
self.ListModels = ListModels
#for m in self.ListModels:
# m.freeze()
input_p = len(self.ListModels)
self.classifier = torch.nn.Sequential(
torch.nn.Linear(in_features=input_p, out_features=128),
torch.nn.SiLU(),
#torch.nn.Dropout(0.1), #0.5
torch.nn.Linear(128, 1),
torch.nn.Sigmoid()
)
self.lossFN = torch.nn.CrossEntropyLoss(label_smoothing=0.5) if typeLossFN == TypeLossFunction.CrossEntropy else torch.nn.BCELoss(reduction="sum")
self.InitWeight = InitWeight
self.scheduler = schedulerLR
self.lr = lr
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1Score = torchmetrics.F1Score('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
def forward(self, data):
data = data.float()
p = [m(data) for m in self.ListModels]
p = torch.cat(p, dim=1)
x = self.classifier(p)
return x
def _shared_step(self, batch):
img, label = batch
if self.InitWeight is not None:
self.lossFN.weight = self.InitWeight[label]
label = label.float()
prediction = self(img)[:,0] #.flatten()
loss = self.lossFN(prediction, label)
predicted_labels = prediction
return loss, label, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.f1Score.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_f1", self.f1Score, on_epoch=True, on_step=False, prog_bar=True)
return loss
def test_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
if self.scheduler:
#scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", patience=1, verbose=True, factor=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 6, 1e-8, verbose=True)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
#"monitor": "val_loss",
}
}
else:
return [optimizer]
class RARP_NVB_ResNet18(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_DaVit(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = timm.create_model("davit_small.msft_in1k", pretrained=True, num_classes=1)
class RARP_NVB_EfficientNetV2_Deep(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.efficientnet_v2_s(weights=torchvision.models.EfficientNet_V2_S_Weights.DEFAULT)
tempFC_ft = self.model.classifier[1].in_features
self.model.classifier[1] = torch.nn.Linear(in_features=tempFC_ft, out_features=128)
self.model.classifier.append(torch.nn.SiLU(True))
self.model.classifier.append(torch.nn.Linear(128, 8))
self.model.classifier.append(torch.nn.SiLU(True))
self.model.classifier.append(torch.nn.Linear(8, 1))
class RARP_NVB_ResNet50_Deep(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Sequential(
torch.nn.Dropout(0.2),
torch.nn.Linear(in_features=tempFC_ft, out_features=128),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8),
torch.nn.SiLU(True),
torch.nn.Linear(8, 1)
)
def forward(self, img):
img = img.float()
pred = self.model(img)
return pred
class RARP_NVB_ResNet50_Deep_OPTuna(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, output_dims=[128, 8], dropuot=0.2, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
layers = []
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
inputDim = self.model.fc.in_features
for outputDim in output_dims:
layers.append(torch.nn.Linear(inputDim, outputDim))
layers.append(torch.nn.SiLU(True))
layers.append(torch.nn.Dropout(dropuot))
inputDim = outputDim
layers.append(torch.nn.Linear(inputDim, 1))
self.model.fc = torch.nn.Sequential(*layers)
def forward(self, img):
img = img.float()
pred = self.model(img)
return pred
class RARP_NVB_ResNet50_V2(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, InputNeurons:int = 8, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Sequential(
torch.nn.Dropout(0.2),
torch.nn.Linear(in_features=tempFC_ft, out_features=128),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8)
)
self.model.fc2 = torch.nn.Linear(8 + InputNeurons, 1)
def forward(self, data):
img, extra = data
img = img.float()
x = torch.nn.functional.silu(self.model(img), True)
extradata = torch.concat((x, extra), dim=1)
pred = self.model.fc2(extradata)
return pred
class RARP_NVB_ResNet50_V3(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, InputNeurons:int = 8, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=128)
self.model.extraFC = torch.nn.Linear(InputNeurons, 128)
self.model.fc2 = torch.nn.Linear(256, 1)
def forward(self, data):
img, extra = data
img = img.float()
x = torch.nn.functional.silu(self.model(img), True)
y = torch.nn.functional.silu(self.model.extraFC(extra), True)
x = torch.concat((x, y), dim=1)
pred = self.model.fc2(x)
return pred
class RARP_NVB_ResNet50_V1(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, InputNeurons:int = 8, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features + InputNeurons
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.Decoder = torch.nn.Sequential(*list(self.model.children())[:-2])
def forward(self, data):
img, extra = data
img = img.float()
featureMap = self.Decoder(img)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(input=featureMap, output_size=(1, 1))
Cont_Net = torch.flatten(Cont_Net, 1)
Cont_Net = torch.concat((Cont_Net, extra), dim=1)
pred = self.model.fc(Cont_Net)
return pred
class RARP_NVB_VAN(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = van.van_b2(pretrained = True)
tempFC_ft = self.model.head.in_features
self.model.head = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_ResNet50(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_MLP(torch.nn.Module):
def __init__(self, in_dim, out_dim, hidden_dim=512, bottleneck=256, n_layers=3, norm_last_layer=False, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.activationFN = torch.nn.GELU()
if n_layers == 1:
self.mlp = torch.nn.Linear(in_dim, bottleneck)
else:
layers = [torch.nn.Linear(in_dim, hidden_dim)]
#layers.append(torch.nn.BatchNorm1d(hidden_dim)) # Add
layers.append(self.activationFN)
layers.append(torch.nn.Dropout(0.30))
for _ in range(n_layers - 2):
layers.append(torch.nn.Linear(hidden_dim, hidden_dim))
layers.append(self.activationFN)
layers.append(torch.nn.Linear(hidden_dim, bottleneck))
self.mlp = torch.nn.Sequential(*layers)
self.apply(self._init_weights)
self.last_layer = torch.nn.utils.weight_norm(
torch.nn.Linear(bottleneck, out_dim, bias=False)
)
self.last_layer.weight_g.data.fill_(1)
if norm_last_layer:
self.last_layer.weight_g.requires_grad = False
def _init_weights(self, m):
if isinstance(m, torch.nn.Linear):
torch.nn.init.normal_(m.weight, std=0.02)
if m.bias is not None:
torch.nn.init.constant_(m.bias, 0)
def forward(self, x):
x = self.mlp(x)
x = torch.nn.functional.normalize(x, dim=-1, p=2)
x = self.last_layer(x)
return x
class RARP_NVB_DINO_Wrapper(torch.nn.Module):
def __init__(self, backbone:torch.nn.Module, new_head:torch.nn.Module, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.backbone = backbone
self.head = new_head
def forward(self, x):
if isinstance(x, list):
n_crops = len(x)
concatCrops = torch.cat(x, dim=0)
else:
concatCrops = x
n_crops = 1
embedding = self.backbone(concatCrops)
logitis = self.head(embedding)
chunks = logitis.chunk(n_crops)
return chunks
class RARP_NVB_DINO_Loss(torch.nn.Module):
def __init__(self, out_dim:int, teacher_Thao:float = 0.04, student_Thao:float = 0.1, center_momentum:float = 0.9, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.S_Thao = student_Thao
self.T_Thao = teacher_Thao
self.C_Momentum = center_momentum
self.register_buffer("center", torch.zeros(1, out_dim))
def forward(self, s_Output, t_Output):
sTemp = [s / self.S_Thao for s in s_Output]
tTemp = [(t - self.center) / self.T_Thao for t in t_Output]
studentSM = [torch.nn.functional.log_softmax(s, dim=-1) for s in sTemp]
teacherSM = [torch.nn.functional.softmax(t, dim=-1).detach() for t in tTemp]
total_loss = 0
n_loss_terms = 0
for t_ix, t in enumerate(teacherSM):
for s_ix, s in enumerate(studentSM):
if (t_ix == s_ix) and (len(teacherSM) > 1):
continue
loss = torch.sum(-t * s, dim=-1)
total_loss += loss.mean()
n_loss_terms += 1
total_loss /= n_loss_terms
self.update_center(t_Output)
return total_loss
@torch.no_grad()
def update_center(self, t_output):
b = torch.cat(t_output).mean(dim=0, keepdim=True)
self.center = self.center * self.C_Momentum + b * (1 - self.C_Momentum)
class RARP_NVB_DINO_RestNet50_Deep(L.LightningModule):
def __init__(
self,
PseudoEstimator: str = None,
threshold: float = 0.5,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher:float = 0.9995,
lr:float = 1e-4,
L1:float = None,
L2:float = 0,
) -> None:
super().__init__()
self.lr = lr
self.Lambda_L1 = L1
self.Lambda_L2 = L2
self.threshold = threshold
self.momentum_teacher = momentum_teacher
self.out_dim = 512
self.in_dim = 2048
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
self.teacher_Labels = RARP_NVB_ResNet50_Deep.load_from_checkpoint(PseudoEstimator, strict=False) if PseudoEstimator is not None else RARP_NVB_ResNet50_Deep()
self.student = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT) #torch.hub.load('facebookresearch/dino:main', 'dino_resnet50')
self.teacher_Features = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
self.student.fc = torch.nn.Identity()
self.teacher_Features.fc = torch.nn.Identity()
self.student = RARP_NVB_DINO_Wrapper(
self.student,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
self.teacher_Features = RARP_NVB_DINO_Wrapper(
self.teacher_Features,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
for parms in self.teacher_Labels.model.parameters():
parms.requires_grad = False
for parms in self.teacher_Features.parameters():
parms.requires_grad = False
self.lossFN_DINO = RARP_NVB_DINO_Loss(self.out_dim, 0.04, 0.1, momentum_teacher)
self.lossFN_KD = torch.nn.CrossEntropyLoss(label_smoothing=0.5) if TypeLoss == TypeLossFunction.CrossEntropy else torch.nn.BCEWithLogitsLoss()
#self.lossFH_KLDiv = torch.nn.KLDivLoss(reduction="batchmean")
self.clasiffier = torch.nn.Sequential(
torch.nn.Dropout(0.2),
torch.nn.Linear(self.out_dim, 128),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8),
torch.nn.SiLU(True),
torch.nn.Linear(8, 1)
)
def forward(self, data, val_step:bool = True):
if val_step:
data = data.float()
dataClassificator, dataTeacher, dataStudent = data, data, data
else:
data = [d.float() for d in data]
dataClassificator, dataTeacher, dataStudent = data[0], data[1:3], data
TeacherDino = self.teacher_Features(dataTeacher)
TeacherLabels = self.teacher_Labels(dataClassificator)
Student = self.student(dataStudent)
# es se evaluan todas las salidas del estuidaitne
#if isinstance(dataStudent, list):
# #index = np.random.randint(0, len(dataStudent))
# temp = self.student(dataStudent)
# CatS_Classifier = torch.cat(temp, dim=0)
# meanS_Classifier = torch.zeros(self.in_dim)
# for dataS in temp:
# meanS_Classifier += dataS
# #S_Classifier = self.student(dataStudent[index])
# S_Classifier = meanS_Classifier / len(dataStudent)
#else:
# S_Classifier = Student
Cont_Net = torch.cat(Student, dim=0)
pred = self.clasiffier(Cont_Net)
if not val_step:
TeacherLabels = [self.teacher_Labels(dataClassificator) for _ in range(len(dataStudent))]
TeacherLabels = torch.cat(TeacherLabels, dim=0)
TeacherLabelsPred = torch.sigmoid(TeacherLabels.flatten())
PseudoLabels = (TeacherLabelsPred > self.threshold) * 1.0
return (pred.flatten(), PseudoLabels, TeacherLabels.flatten()), (TeacherDino, Student)
def _shared_step(self, batch, val_step:bool = False):
img, label = batch
if not val_step:
label = torch.cat([label for _ in range(len(img))], dim=0)
label = label.float()
KD_Prediction, DINO_Loss = self(img, val_step)
TeacherF, StudentF = DINO_Loss
prediction, PseudoLabels, teacherOutputs = KD_Prediction
predicted_labels = torch.sigmoid(prediction)
##verstion 1
W_Alpha, W_Beta = (1, 0.5)#(1, 0.5)
loss = W_Alpha * self.lossFN_KD(prediction, PseudoLabels) + W_Beta * self.lossFN_KD(prediction, label)
#version 2
#thao_KD = 1#5.0
#W_Alpha, W_Beta = (0.6, 0.4)
#softTeacher = torch.sigmoid(teacherOutputs/thao_KD)
#softStudent = torch.sigmoid(prediction/thao_KD)
#loss_sl = torch.nn.functional.binary_cross_entropy(softStudent, softTeacher)
#loss_hl = self.lossFN_KD(prediction, label)
#loss = W_Alpha * loss_hl + W_Beta * loss_sl
#loss = W_Alpha * self.lossFN_KD(prediction, label) + W_Beta * (self.lossFH_KLDiv(softStudent, softTeacher) * (thao_KD ** 2))
loss += (self.lossFN_DINO(StudentF, TeacherF) if not val_step else 0)
if not val_step:
self.logger.experiment.add_histogram ("Teacher", TeacherF[0])
self.logger.experiment.add_histogram ("Student", StudentF[1])
if self.Lambda_L1 is not None:
loss_l1 = 0
for params in self.student.parameters(): # aqui
loss_l1 += torch.sum(torch.abs(params))
loss += self.Lambda_L1 * loss_l1
#return loss, PseudoLabels, predicted_labels
return loss, label, predicted_labels
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch, False)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def on_train_batch_end(self, outputs, batch, batch_idx):
with torch.no_grad():
for student_ps, teacher_ps in zip(self.student.parameters(), self.teacher_Features.parameters()):
teacher_ps.data.mul_(self.momentum_teacher)
teacher_ps.data.add_((1-self.momentum_teacher) * student_ps.detach().data)
self.logger.experiment.add_histogram ("Teacher_Center", self.lossFN_DINO.center)
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels = self._shared_step(batch, True)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.student.parameters(), lr=self.lr, weight_decay=self.Lambda_L2)
return [optimizer]
class RARP_NVB_DINO_VAN(RARP_NVB_DINO_RestNet50_Deep):
def __init__(
self,
PseudoEstimator: str = None,
threshold: float = 0.5,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher: float = 0.9995,
lr: float = 0.0001,
L1: float = None,
L2: float = 0
) -> None:
super().__init__(PseudoEstimator, threshold, TypeLoss, momentum_teacher, lr, L1, L2)
self.in_dim = 512
self.student = van.van_b2(pretrained = True, num_classes = -1)
self.teacher_Features = van.van_b2(pretrained = True, num_classes = -1)
self.student = RARP_NVB_DINO_Wrapper(
self.student,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
self.teacher_Features = RARP_NVB_DINO_Wrapper(
self.teacher_Features,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
for parms in self.teacher_Features.parameters():
parms.requires_grad = False
class RARP_NVB_DINO_ViT(RARP_NVB_DINO_RestNet50_Deep):
def __init__(
self,
PseudoEstimator: str = None,
threshold: float = 0.5,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher: float = 0.9995,
lr: float = 0.0001,
L1: float = None,
L2: float = 0
) -> None:
super().__init__(PseudoEstimator, threshold, TypeLoss, momentum_teacher, lr, L1, L2)
self.in_dim = 768
self.student = torchvision.models.vit_b_16(weights=torchvision.models.ViT_B_16_Weights.DEFAULT)
self.teacher_Features = torchvision.models.vit_b_16(weights=torchvision.models.ViT_B_16_Weights.DEFAULT)
self.student.heads = torch.nn.Identity()
self.teacher_Features.heads = torch.nn.Identity()
self.student = RARP_NVB_DINO_Wrapper(
self.student,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
self.teacher_Features = RARP_NVB_DINO_Wrapper(
self.teacher_Features,
RARP_NVB_MLP(self.in_dim, self.out_dim, 1024)
)
for parms in self.teacher_Features.parameters():
parms.requires_grad = False
class RARP_MAE(L.LightningModule):
def __init__(
self,
backbone:str,
mask_ratio:float = 0.75,
patch_size: int = 16,
#img_size: int = 224,
lr: float = 1e-4,
hiden_channels = [512, 256, 128, 64, 32],
img_mean_std = [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
activation_fn = torch.nn.ReLU(True)
):
super().__init__()
self.save_hyperparameters(ignore=["activation_fn", "img_mean_std"])
if img_mean_std is not None:
mean, std = img_mean_std
self.register_buffer("mean_IMG", torch.tensor(mean).view(3, 1, 1))
self.register_buffer("std_IMG", torch.tensor(std).view(3, 1, 1))
else:
self.mean_IMG = None
self.std_IMG = None
self.loss_fn = torch.nn.MSELoss()
match backbone:
case "resnet":
model = torchvision.models.resnet18()
model.fc = torch.nn.Identity()
self.encoder = model
self.encoder_out_dim = 512
case "van":
model = van.van_b2()
model.head = torch.nn.Identity()
self.encoder = model
self.encoder_out_dim = 512
case "levit":
self.encoder = timm.create_model("levit_192.fb_dist_in1k", pretrained=False, num_classes=0)
self.encoder_out_dim = 384
hiden_channels = [384, 256, 128, 64, 32]
in_channel = self.encoder_out_dim
layers = [
torch.nn.Linear(in_channel, hiden_channels[0]*7*7),
torch.nn.Unflatten(1, (hiden_channels[0], 7, 7)),
]
for out_channel in hiden_channels[1:]:
layers.append(torch.nn.ConvTranspose2d(in_channel, out_channel, kernel_size=3, stride=2, padding=1, output_padding=1))
layers.append(torch.nn.BatchNorm2d(out_channel))
layers.append(activation_fn)
in_channel = out_channel
#Last layer
layers.append(torch.nn.ConvTranspose2d(in_channel, 3, kernel_size=3, stride=2, padding=1, output_padding=1))
self.decoder = torch.nn.Sequential(*layers)
def _denormalize(self, tensor:torch.Tensor):
return tensor * self.std_IMG + self.mean_IMG
def _mask_patches(self, imgs):
B, C, H, W = imgs.shape
ph, pw = self.hparams.patch_size, self.hparams.patch_size
gh, gw = H // ph, W // pw
mask:torch.Tensor = (torch.rand(B, gh * gw, device=imgs.device) >= self.hparams.mask_ratio)
mask = mask.reshape(B, 1, gh, gw) # [B,1,gh,gw]
# expand mask to full image
mask = mask.repeat_interleave(ph, dim=2) # [B,1,gh*ph,gw]
mask = mask.repeat_interleave(pw, dim=3) # [B,1,gh*ph,gw*pw] == [B,1,H,W]
mask = mask.expand(-1, C, -1, -1)
imgs_masked = imgs * mask
return imgs_masked
def forward(self, data, val_step=False):
if not val_step:
data = [d.float() for d in data] #[0] original Img, [1] augmented Img
else:
data = [data.float(), data.float()]
imgs_masked = self._mask_patches(data[1])
latent = self.encoder(imgs_masked) # [B, 512]
reconstructed_image = self.decoder(latent) # [B, 3, 224, 224]
return reconstructed_image, data[0]
def _shared_step(self, batch, val_step=False):
img, _ = batch
res, true_img = self(img, val_step)
loss_MSE = self.loss_fn(res, true_img)
return loss_MSE, res, true_img
def training_step(self, batch, batch_idx):
loss, img, _ = self._shared_step(batch, False)
self.log("train_loss", loss, on_epoch=True, sync_dist=True)
if batch_idx % 500 == 0 and self.mean_IMG is not None and self.std_IMG is not None:
imgReconstruction = torch.clip(self._denormalize(img), 0, 1)
grid = torchvision.utils.make_grid(imgReconstruction)
self.logger.experiment.add_image('reconstructed_images', grid, self.global_step)
return loss
def validation_step(self, batch, batch_idx):
loss, img, target_img = self._shared_step(batch, True)
self.log("val_loss", loss, on_epoch=True, prog_bar=True)
if self.mean_IMG is not None and self.std_IMG is not None:
img = torch.clip(self._denormalize(img), 0, 1)
target_img = torch.clip(self._denormalize(target_img), 0, 1)
ssim = piq.ssim(img, target_img, data_range=1.0)
psnr = piq.psnr(img, target_img, data_range=1.0)
self.log("val_ssim", ssim, on_epoch=True)
self.log("val_psnr", psnr, on_epoch=True)
if batch_idx % 100 == 0:
grid = torchvision.utils.make_grid(img)
self.logger.experiment.add_image('val_reconstructed_images', grid, self.global_step)
def on_after_backward(self):
total_norm = 0.0
for p in self.parameters():
if p.grad is not None:
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
self.log("grad_norm", total_norm)
if total_norm < 1e-8:
self.log("grad_warning", "Vanishing gradient suspected!")
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)
return [optimizer]
class RARP_Encoder_DINO(L.LightningModule):
def __init__(self,
momentum_teacher:float = 0.9995,
lr:float = 1e-4,
Teacher_T:float = 0.04,
Student_T:float = 0.1,
max_epochs:int = 100,
total_steps:int = None
) -> None:
super().__init__()
self.save_hyperparameters()
self.lr = lr
self.momentum_teacher = momentum_teacher
self.out_dim = 65536
self.in_dim = 768 #512
self.embs_t = []
self.embs_s = []
#self.student = van.van_b2(num_classes = 0)
#self.teacher = van.van_b2(num_classes = 0)
#weights=torchvision.models.ConvNeXt_Small_Weights.DEFAULT
self.student = torchvision.models.convnext_small()
self.student.classifier[-1] = torch.nn.Identity()
self.teacher = torchvision.models.convnext_small()
self.teacher.classifier[-1] = torch.nn.Identity()
self.student = RARP_NVB_DINO_Wrapper(
self.student,
RARP_NVB_MLP(self.in_dim, self.out_dim, hidden_dim=2048, bottleneck=256, norm_last_layer=True)
)
self.teacher = RARP_NVB_DINO_Wrapper(
self.teacher,
RARP_NVB_MLP(self.in_dim, self.out_dim, hidden_dim=2048, bottleneck=256, norm_last_layer=True)
)
self.teacher.load_state_dict(self.student.state_dict())
for parms in self.teacher.parameters():
parms.requires_grad = False
self.lossFN_DINO = RARP_NVB_DINO_Loss(self.out_dim, Teacher_T, Student_T, momentum_teacher)
def forward(self, data):
data = [d.float() for d in data]
dataTeacher, dataStudent = data[0:3], data
teacher_features = self.teacher(dataTeacher)
student_features = self.student(dataStudent)
return teacher_features, student_features
def _shared_step(self, batch):
img, _ = batch
t, s = self(img)
loss_Dino = self.lossFN_DINO(s, t)
return loss_Dino
def training_step(self, batch, batch_idx):
loss = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True, sync_dist=True)
return loss
def on_train_batch_end(self, outputs, batch, batch_idx):
#step = self.global_step
#
#m = 1.0 - (1.0 - self.hparams.momentum_teacher) * (
# (1 + math.cos(math.pi * step / self.hparamstotal_steps)) / 2
#)
with torch.no_grad():
for student_ps, teacher_ps in zip(self.student.parameters(), self.teacher.parameters()):
teacher_ps.data.mul_(self.momentum_teacher)
teacher_ps.data.add_((1-self.momentum_teacher) * student_ps.detach().data)
def validation_step(self, batch, batch_idx):
imgs, _ = batch
embs = self.teacher(imgs.float())[0]
self.embs_t.append(embs)
embs = self.student(imgs.float())[0]
self.embs_s.append(embs)
def on_validation_epoch_end(self):
emds = torch.cat(self.embs_t, dim=0).cpu().numpy()
self.embs_t.clear()
#silhouette
kmeans = KMeans(n_clusters=10, random_state=505).fit(emds)
sil = silhouette_score(emds, kmeans.labels_)
self.log("val_silhouette_teacher", sil)
emds = torch.cat(self.embs_s, dim=0).cpu().numpy()
self.embs_s.clear()
#silhouette
kmeans = KMeans(n_clusters=10, random_state=505).fit(emds)
sil = silhouette_score(emds, kmeans.labels_)
self.log("val_silhouette_student", sil)
def on_after_backward(self):
norms = [p.grad.data.norm(2).item() for p in self.parameters() if p.grad is not None]
avg_layer_norm = sum(norms) / len(norms)
self.log("avg_grad_norm", avg_layer_norm)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.student.parameters(), lr=self.lr)
return [optimizer]
class RARP_NVB_Classification_Head(torch.nn.Module):
def __init__(self, in_features:int, out_features:int, layer:list=[], activation_fn:torch.nn.Module = torch.nn.ReLU(), *args, **kwargs):
super().__init__(*args, **kwargs)
self.activation = activation_fn
if len (layer) == 0:
self.head = torch.nn.Linear(in_features, out_features)
else:
temp_head = []
next_input = in_features
for num in layer:
temp_head.append(torch.nn.Linear(next_input, num))
temp_head.append(self.activation)
temp_head.append(torch.nn.Dropout(0.4))
next_input = num
temp_head[-1] = torch.nn.Dropout(0.2)
temp_head.append(torch.nn.Linear(next_input, out_features))
self.head = torch.nn.Sequential(*temp_head)
del temp_head
def forward(self, x):
return self.head(x)
class RARP_Encoder_DINO_AUX_task(RARP_Encoder_DINO):
def __init__(
self,
momentum_teacher = 0.9995,
lr = 0.0001,
Teacher_T = 0.04,
Student_T = 0.1,
max_epochs = 100,
total_steps = None,
aux_lambda = 0.3
):
super().__init__(momentum_teacher, lr, Teacher_T, Student_T, max_epochs, total_steps)
self.angles = [-90, 0, 90, 180]
self.angles_labels = torch.from_numpy(LabelEncoder().fit_transform(self.angles))
self.classifier = RARP_NVB_Classification_Head(self.in_dim, len(self.angles), layer=[128], activation_fn=torch.nn.GELU())
self.lossFN_Aux = torch.nn.CrossEntropyLoss()
self.train_acc = torchmetrics.Accuracy("multiclass", num_classes=len(self.angles))
self.val_acc = torchmetrics.Accuracy("multiclass", num_classes=len(self.angles))
def _rotation_labels_generator(self, batch_size:int, angles_batch=None, rand:bool=False) -> torch.Tensor:
assert len(angles_batch) > 0, "Empty list, angles list shuld have more than 2 values"
if rand:
return torch.tensor([angles_batch[np.random.randint(len(angles_batch))] for _ in range(batch_size)])
else:
return torch.tensor([angles_batch[i % len(angles_batch)] for i in range(batch_size)])
def _rotate_img(self, imgs:torch.Tensor, angles_batch=[]) -> torch.Tensor:
assert len(angles_batch) > 0, "Empty list, angles list shuld have more than 2 values"
list_imgs = []
for i, angle_idx in enumerate(angles_batch):
list_imgs.append(torchvision.transforms.functional.rotate(imgs[i, ...], self.angles[angle_idx]))
return torch.stack(list_imgs, dim=0)
def forward(self, data, rot_data, val_step:bool=False):
if not val_step:
data = [d.float() for d in data]
dataTeacher, dataStudent = data[0:3], data
else:
data = data.float()
dataTeacher, dataStudent = data, data
dataAux = rot_data.float()
teacher_features = self.teacher(dataTeacher)
student_features = self.student(dataStudent)
aux_rot = self.student.backbone(dataAux)
aux_rot = self.classifier(aux_rot)
return teacher_features, student_features, aux_rot
def _shared_step(self, batch, val_step:bool=False):
img, _ = batch
batch_size = img.size(0) if val_step else img[0].size(0)
current_device = img.device if val_step else img[0].device
labels = self._rotation_labels_generator(batch_size, self.angles_labels, not val_step).to(current_device)
rot_img = self._rotate_img(img if val_step else img[0], labels).to(current_device)
t, s, aux = self(img, rot_img, val_step)
if val_step:
self.embs_t.append(t[0])
self.embs_s.append(s[0])
aux = torch.softmax(aux, dim=1)
loss_aux = self.lossFN_Aux(aux, labels)
loss_aux = self.hparams.aux_lambda * loss_aux
loss_Dino = self.lossFN_DINO(s, t) if not val_step else torch.tensor(0, device=current_device, dtype=torch.float32)
loss = loss_Dino + loss_aux
return loss, labels, aux, (loss_Dino, loss_aux)
def training_step(self, batch, batch_idx):
loss, true_labels, pred_labels, losses = self._shared_step(batch, False)
self.log("train_loss", loss, on_epoch=True, sync_dist=True)
self.train_acc.update(pred_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
self.log("train_loss_DINO", losses[0], on_epoch=True, on_step=False)
self.log("train_loss_AUX", losses[1], on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, pred_labels, losses = self._shared_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(pred_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_loss_DINO", losses[0], on_epoch=True, on_step=False)
self.log("val_loss_AUX", losses[1], on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.AdamW([
{"params": self.student.parameters(), "lr": self.lr},
{"params": self.classifier.parameters(), "lr": self.lr}
])
return [optimizer]
class RARP_NVB_DINO_MultiTask(L.LightningModule):
# Define a hook function to capture the output
def _hook_fn_Student(self, module, input, output):
self.last_conv_output_S = output
def _hook_fn_Teacher(self, module, input, output):
self.last_conv_output_T = output
def __init__(
self,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher:float = 0.9995,
lr:float = 1e-4,
L1:float = None,
L2:float = 0,
std: float = None,
mean: float = None,
SoftAdptAlgo:int = 0,
SoftAdptBeta:float = 0.1,
Teacher_T:float = 0.04,
Student_T:float = 0.1,
intermittent:bool = False
) -> None:
super().__init__()
self.intermittent_train = intermittent
self.std_IMG = torch.tensor(std).view(3, 1, 1) if std is not None else None
self.mean_IMG = torch.tensor(mean).view(3, 1, 1) if mean is not None else None
self.lr = lr
self.Lambda_L1 = L1
self.Lambda_L2 = L2
self.Teacher_t = Teacher_T
self.Studet_T = Student_T
self.momentum_teacher = momentum_teacher
self.out_dim = 1024
self.in_dim = 512
self.weights = torch.tensor([1,1,1])
self.softAdapt = NormalizedSoftAdapt(SoftAdptBeta) if SoftAdptAlgo == 1 else LossWeightedSoftAdapt(SoftAdptBeta)
self.loss_history = {
'loss_DINO': [],
'loss_Reconstruction': [],
'loss_Binary': [],
}
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
self.student = van.van_b2(pretrained = True, num_classes = 0)
self.teacher_Features = van.van_b2(pretrained = True, num_classes = 0)
self.decoder = DynamicDecoder(input_channels=1024)
self.student = RARP_NVB_DINO_Wrapper(
self.student,
RARP_NVB_MLP(self.in_dim, self.out_dim, n_layers=2)
)
self.teacher_Features = RARP_NVB_DINO_Wrapper(
self.teacher_Features,
RARP_NVB_MLP(self.in_dim, self.out_dim, n_layers=2)
)
self.teacher_Features.load_state_dict(self.student.state_dict())
for parms in self.teacher_Features.parameters():
parms.requires_grad = False
self.lossFN_DINO = RARP_NVB_DINO_Loss(self.out_dim, Teacher_T, Student_T, momentum_teacher)
self.lossFN = torch.nn.CrossEntropyLoss(label_smoothing=0.5) if TypeLoss == TypeLossFunction.CrossEntropy else torch.nn.BCEWithLogitsLoss()
self.ReconstructionLoss = torch.nn.MSELoss()
self.last_conv_output_T = None
self.last_conv_output_S = None
self.teacher_Features.backbone.block4[-1].register_forward_hook(self._hook_fn_Teacher)
self.student.backbone.block4[-1].register_forward_hook(self._hook_fn_Student)
self.clasiffier = torch.nn.Sequential(
torch.nn.Linear(1024, 128),
torch.nn.SiLU(True),
torch.nn.Dropout(0.4),
torch.nn.Linear(128, 8),
torch.nn.SiLU(True),
torch.nn.Dropout(0.2),
torch.nn.Linear(8, 1)
)
print(f"lr={self.lr}, L1={self.Lambda_L1}")
def _denormalize(self, tensor:torch.Tensor):
# Move mean and std to the same device as the input tensor
mean = self.mean_IMG.to(tensor.device)
std = self.std_IMG.to(tensor.device)
return tensor * std + mean
def _calc_L1(self, params):
l1 = 0
for p in params:
l1 += torch.sum(torch.abs(p))
return self.Lambda_L1 * l1
def _calc_weights(self, log_weights:bool = True):
self.weights = self.softAdapt.get_component_weights(
torch.tensor(self.loss_history["loss_DINO"][:-1] if len(self.loss_history["loss_DINO"]) % 2 == 0 else self.loss_history["loss_DINO"]),
torch.tensor(self.loss_history["loss_Reconstruction"][:-1] if len(self.loss_history["loss_Reconstruction"]) % 2 == 0 else self.loss_history["loss_Reconstruction"]),
torch.tensor(self.loss_history["loss_Binary"][:-1] if len(self.loss_history["loss_Binary"]) % 2 == 0 else self.loss_history["loss_Binary"]),
verbose=False
)
if log_weights:
self.log("W_loss_img", self.weights[1], on_epoch=True, on_step=False)
self.log("W_loss_DINO", self.weights[0], on_epoch=True, on_step=False)
self.log("W_loss_GT", self.weights[2], on_epoch=True, on_step=False)
self.loss_history = {
'loss_DINO': [],
'loss_Reconstruction': [],
'loss_Binary': [],
}
def forward(self, data, val_step:bool = True):
if val_step:
data = data.float()
dataTeacher, dataStudent = data, data
else:
data = [d.float() for d in data]
dataTeacher, dataStudent = data[1:3], data
TeacherDino = self.teacher_Features(dataTeacher)
Student = self.student(dataStudent)
if not val_step:
NumChunks = len(dataStudent)
S_GlogalViews = self.last_conv_output_S.chunk(NumChunks)[1:3]
self.last_conv_output_S = torch.cat(S_GlogalViews, dim=0)
cat_features = torch.cat((self.last_conv_output_S, self.last_conv_output_T), dim=1)
reconstructed_image = self.decoder(cat_features)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(cat_features, (1,1)).flatten(1)
pred = self.clasiffier(Cont_Net)
return pred, (Student, TeacherDino), reconstructed_image
def _shared_step(self, batch, val_step:bool = False):
img, label = batch
prediction, features, new_image = self(img, val_step)
StudentF, TeacherF = features
if isinstance(self.clasiffier, torch.nn.Sequential):
if self.clasiffier[-1].out_features == 1:
prediction = prediction.flatten()
elif isinstance(self.clasiffier, (NOAH, RARP_NVB_Classification_Head)):
prediction = prediction.flatten()
predicted_labels = torch.sigmoid(prediction)
orignalImg = torch.cat([img[0].float() for _ in range(len(TeacherF))], dim=0) if not val_step else img.float()
label = torch.cat([label.float() for _ in range(len(TeacherF))], dim=0) if not val_step else label.float()
#DINO Loss
loss_Dino = self.lossFN_DINO(StudentF, TeacherF) if not val_step else torch.tensor(0, device=label.device, dtype=torch.float32)
#Clasificator
loss_HL = self.lossFN(prediction, label)
#Reconstruction
loss_img = self.ReconstructionLoss(new_image, orignalImg)
loss_img = loss_img.float()
if not val_step:
if self.Lambda_L1 is not None:
loss_HL += self._calc_L1(self.clasiffier.parameters())
if self.Lambda_L2 > 0:
l2_reg = 0.0
for param in self.clasiffier.parameters():
l2_reg += torch.norm(param, 2) ** 2
loss_HL += self.Lambda_L2 * l2_reg
self.loss_history["loss_DINO"].append(loss_Dino.item())
self.loss_history["loss_Reconstruction"].append(loss_img.item())
self.loss_history["loss_Binary"].append(loss_HL.item())
loss = self.weights[0] * loss_Dino + self.weights[1] * loss_img + self.weights[2] * loss_HL
return loss, label, predicted_labels, (self.weights[0] * loss_Dino, self.weights[2] * loss_HL, self.weights[1] * loss_img, new_image)
def on_train_epoch_start(self):
if self.current_epoch % 2 == 0 and self.current_epoch != 0:
self._calc_weights()
if self.intermittent_train and self.current_epoch != 0:
par_epoch = (self.current_epoch % 2 == 0)
for parms in self.student.backbone.parameters():
parms.requires_grad = par_epoch
for parms in self.decoder.parameters():
parms.requires_grad = not par_epoch
for parms in self.clasiffier.parameters():
parms.requires_grad = not par_epoch
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch, False)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
self.log("train_loss_img", losses[2], on_epoch=True, on_step=False)
self.log("train_loss_DINO", losses[0], on_epoch=True, on_step=False)
self.log("train_loss_GT", losses[1], on_epoch=True, on_step=False)
if batch_idx % 50 == 0 and self.mean_IMG is not None and self.std_IMG is not None:
imgReconstruction = torch.clip(self._denormalize(losses[3]) / 255, 0, 1)
imgReconstruction = imgReconstruction[:, [2, 1, 0], :, :]
grid = torchvision.utils.make_grid(imgReconstruction)
self.logger.experiment.add_image('reconstructed_images', grid, self.global_step)
return loss
def on_train_batch_end(self, outputs, batch, batch_idx):
with torch.no_grad():
for student_ps, teacher_ps in zip(self.student.parameters(), self.teacher_Features.parameters()):
teacher_ps.data.mul_(self.momentum_teacher)
teacher_ps.data.add_((1-self.momentum_teacher) * student_ps.detach().data)
#self.logger.experiment.add_histogram ("Teacher_Center", self.lossFN_DINO.center)
def on_after_backward(self):
total_norm = 0.0
for p in self.parameters():
if p.grad is not None:
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
self.log("grad_norm", total_norm)
if total_norm < 1e-8:
self.log("grad_warning", "Vanishing gradient suspected!")
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_loss_img", losses[2], on_epoch=True, on_step=False)
self.log("val_loss_DINO", losses[0], on_epoch=True, on_step=False)
self.log("val_loss_GT", losses[1], on_epoch=True, on_step=False)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels, losses = self._shared_step(batch, True)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
if self.mean_IMG is not None and self.std_IMG is not None:
imgReconstruction = torch.clip(self._denormalize(losses[3]) / 255, 0, 1)
imgReconstruction = imgReconstruction[:, [2, 1, 0], :, :]
imgOrig = torch.clip(self._denormalize(batch[0])/255, 0, 1)
imgOrig = imgOrig[:, [2, 1, 0], :, :]
imgReconstruction = torch.cat((imgOrig, imgReconstruction), dim=0)
grid = torchvision.utils.make_grid(imgReconstruction)
self.logger.experiment.add_image('reconstructed_images_test', grid, self.global_step)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr) #, weight_decay=self.Lambda_L2
return [optimizer]
class Scalar_AttnPooling(torch.nn.Module):
def __init__(self, hidden_dim=64):
super().__init__()
self.proj = torch.nn.Linear(1, hidden_dim)
self.attn = torch.nn.Linear(hidden_dim, 1)
self.pooler_cls = torch.nn.Linear(1, 1)
def forward(self, logits_bt):
x = logits_bt.unsqueeze(-1) #[B, T, 1]
h = torch.tanh(self.proj(x)) #[B, T, hidden_dim]
a = self.attn(h).squeeze(-1) #[B, T]
w = torch.nn.functional.softmax(a, 1) # [B, T] Weights for each T, to idetnfy the best frame for classification
z = (w.unsqueeze(-1) * x).sum(dim=1) # [B, 1]
z = self.pooler_cls(z)
return z
class Scalar_TCN(torch.nn.Module):
def __init__(self, hidden_dim=64, layers=3):
super().__init__()
self.in_proj = torch.nn.Conv1d(1, hidden_dim, kernel_size=1)
blocks = []
for i in range(layers):
d = 2 ** i
blocks += [
torch.nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=d, dilation=d),
torch.nn.SiLU(),
torch.nn.Dropout(0.3),
torch.nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1),
torch.nn.SiLU(),
torch.nn.Dropout(0.2)
]
self.tcn = torch.nn.Sequential(*blocks)
self.head = torch.nn.Linear(hidden_dim, 1)
def forward (self, logits_bt):
x = logits_bt.unsqueeze(1) #[B, 1, T] to do the Conv over the channel or the logits values
x = self.in_proj(x) #[B, hidden_dim, T]
x = self.tcn(x) #[B, hidden_dim, T]
x = x.mean(dim=2) #[B, hidden_dim] global average over T
x = self.head(x) #[B, 1]
return x
class ModuleWrapper(torch.nn.Module):
def __init__(self, module):
super().__init__()
self.module = module
def forward(self, x, dummy_arg=None):
assert dummy_arg is not None
x = self.module(x)
return x
class Chomp1d(torch.nn.Module):
"""
Remove extra padding at the end to maintain causality.
If you pad (padding) at left, you may need to chomp off the right extra.
"""
def __init__(self, chomp_size):
super().__init__()
self.chomp_size = chomp_size
def forward(self, x):
# x has shape [B, C, T]
if self.chomp_size == 0:
return x
return x[:, :, :-self.chomp_size]
class TemporalBlock(torch.nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, dilation, padding, dropout=0.0):
"""
A residual block in TCN with two dilated conv layers (same dilation).
"""
super().__init__()
self.conv1 = torch.nn.Conv1d(in_channels, out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation)
self.chomp1 = Chomp1d(padding)
self.relu1 = torch.nn.ReLU()
self.dropout1 = torch.nn.Dropout(dropout)
self.conv2 = torch.nn.Conv1d(out_channels, out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation)
self.chomp2 = Chomp1d(padding)
self.relu2 = torch.nn.ReLU()
self.dropout2 = torch.nn.Dropout(dropout)
self.downsample = (torch.nn.Conv1d(in_channels, out_channels, 1)
if in_channels != out_channels else None)
self.relu = torch.nn.ReLU()
def forward(self, x):
"""
x: [B, in_channels, T]
returns: [B, out_channels, T]
"""
out = self.conv1(x)
out = self.chomp1(out)
out = self.relu1(out)
out = self.dropout1(out)
out = self.conv2(out)
out = self.chomp2(out)
out = self.relu2(out)
out = self.dropout2(out)
res = x if self.downsample is None else self.downsample(x)
return self.relu(out + res)
class TemporalConvNet(torch.nn.Module):
def __init__(self, num_inputs, num_channels, kernel_size=3, dropout=0.0):
"""
num_inputs: number of input channels (features)
num_channels: list of output channels per layer, e.g. [64, 64, 128]
kernel_size: convolution kernel size (e.g. 3)
dropout: dropout rate in blocks
"""
super().__init__()
layers = []
num_levels = len(num_channels)
for i in range(num_levels):
in_ch = num_inputs if i == 0 else num_channels[i - 1]
out_ch = num_channels[i]
dilation = 2 ** i
# padding should be such that the output has length T (causal)
padding = (kernel_size - 1) * dilation
layers.append(
TemporalBlock(in_ch, out_ch,
kernel_size=kernel_size,
stride=1,
dilation=dilation,
padding=padding,
dropout=dropout)
)
self.network = torch.nn.Sequential(*layers)
def forward(self, x):
"""
x: [B, T, C_in]
returns: [B, C_out_last, T]
"""
x = x.permute(0, 2, 1) #[B, C_in, T]
x = self.network (x) #[B, C_out, T]
x = x.permute(0, 2, 1) #[B, T, C_out]
return x
class RARP_NVB_DINO_MultiTask_A5_Video(L.LightningModule):
def __init__(
self,
base_model_path = None,
lr = 0.0001,
wd = 0.01,
L1 = None,
L2 = 0,
std = None,
mean = None,
head_type:int = 0, #None = 0, linear = 1, Attn. Pooling = 2, TCN = 3
chunks_loading:int = 50
):
super().__init__()
self.lr = lr
self.wd = wd
self.chunks = chunks_loading
if base_model_path is not None:
self.base_model = RARP_NVB_DINO_MultiTask.load_from_checkpoint(base_model_path)
self.base_model.eval()
for param in self.base_model.parameters():
param.requires_grad = False
else:
self.base_model = RARP_NVB_DINO_MultiTask()
self.lossFN = torch.nn.BCEWithLogitsLoss()
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
#match(head_type):
# case 1:
# #Linear
# self.head = torch.nn.Linear(600, 1)
# case 2:
# #Attn. pooling
# self.head = Scalar_AttnPooling(128)
# case 3:
# #TCN
# self.head = Scalar_TCN(64, 1)
# case _:
self.head = torch.nn.Linear(600, 1)
def _shared_video_step(self, batch:list[torch.Tensor], val_step:bool = False):
video, label = batch
B, T, C, H, W = video.shape
video = video.float() #[B, T, C, H, W]
label = label.float() #[B]
chunk_T = self.chunks
pred_bt = []
def _fn(inp):
pred, *_ = self.base_model(inp)
return pred
for t0 in tqdm(range(0, T, chunk_T), desc=f"Video Analysis in {chunk_T} chunk", leave=False):
t1 = min(T, t0 + chunk_T)
x = video[:, t0:t1].reshape(-1, C, H, W).contiguous(memory_format=torch.channels_last)
pred = torch_ckp.checkpoint(_fn, x)
pred_bt.append(pred.view(B, t1-t0, -1))
pred_video = torch.cat(pred_bt, dim=1).flatten(start_dim=1)
pred_video = self(pred_video, val_step) #[B, 1]
pred_video = pred_video.flatten() #[B] to match labels shape
predicted_labels = torch.sigmoid(pred_video)
loss = self.lossFN(pred_video, label)
return loss, label, predicted_labels
def forward(self, data, val_step:bool = True):
#if self.head is None:
# pred_video = data.mean(dim=1)
#else:
pred_video = self.head(data)
return pred_video
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, False)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.head.parameters(), lr=self.lr, weight_decay=self.wd) #, weight_decay=self.Lambda_L2
return [optimizer]
class RARP_NVB_VIDEO_3D_ResNet(L.LightningModule):
def __init__(
self,
lr = 0.0001,
wd = 0.01,
L1 = None,
L2 = 0,
std = None,
mean = None,
chunks_loading:int = 50,
str_path:str = None
):
super().__init__()
self.lr = lr
self.wd = wd
self.chunks = chunks_loading
if str_path is None:
base_2D_model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
self.check_pt = True
self.base_model = I3DResNet50( base_2D_model )
self.features = self.base_model.fc.in_features
self.base_model.fc = torch.nn.Linear(self.features, 1)
else:
base_2D_model = torchvision.models.resnet50()
in_f = base_2D_model.fc.in_features
base_2D_model.fc = torch.nn.Linear(in_f, 1)
base_2D_model.load_state_dict(torch.load(str_path))
self.check_pt = True
self.base_model = I3DResNet50( base_2D_model )
self.lossFN = torch.nn.BCEWithLogitsLoss()
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
self.test_probs = []
self.test_targets = []
def _shared_video_step(self, batch:list[torch.Tensor], val_step:bool = False):
video, label = batch
video = video.float() #[B, T, C, H, W]
label = label.float() #[B]
pred_video = self(video) #[B, T, 1]
#pred_video = pred_video.mean(dim=1) #[B, 1]
pred_video = pred_video.flatten() #[B] to match labels shape
predicted_labels = torch.sigmoid(pred_video)
loss = self.lossFN(pred_video, label)
return loss, label, predicted_labels
def forward(self, video:torch.Tensor):
T = video.shape[1]
chunk_T = self.chunks
pred_bt = []
assert T > chunk_T, f"The Time dim is smaller, than chunk size [T={T}, chunk={chunk_T}]"
def _fn(inp):
pred = self.base_model(inp)
return pred
video = video.permute(0, 2, 1, 3, 4)
for t0 in tqdm(range(0, T, chunk_T), desc=f"Video Analysis in {chunk_T} chunk", leave=False): # Loop for each chunk
t1 = min(T, t0 + chunk_T)
x = video[:, :, t0:t1]#.contiguous(memory_format=torch.channels_last_3d)
pred = torch_ckp.checkpoint(_fn, x, use_reentrant=False) #froward to CNN and checkpoint grads
pred_bt.append(pred) #[B, 1]
pred_video = torch.stack(pred_bt, dim=0).mean(dim=0)
#pred_video = torch.mean(pred_bt, dim=1) #concat all chunks -> [B, T, C, H, W]
#pred_video = self.head(pred_video)
return pred_video
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.wd) #, weight_decay=self.Lambda_L2
return [optimizer]
def on_after_backward(self):
norms = [p.grad.data.norm(2).item() for p in self.parameters() if p.grad is not None]
avg_layer_norm = sum(norms) / len(norms)
self.log("avg_grad_norm", avg_layer_norm)
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, False)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.test_probs.append(predicted_labels)
self.test_targets.append(true_labels)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def on_test_epoch_end(self):
predicted_labels = torch.cat(self.test_probs).to(self.device)
true_labels = torch.cat(self.test_targets).to(self.device).int()
acc = torchmetrics.Accuracy('binary').to(self.device)(predicted_labels, true_labels)
precision = torchmetrics.Precision('binary').to(self.device)(predicted_labels, true_labels)
recall = torchmetrics.Recall('binary').to(self.device)(predicted_labels, true_labels)
auc = torchmetrics.AUROC('binary').to(self.device)(predicted_labels, true_labels)
f1Score = torchmetrics.F1Score('binary').to(self.device)(predicted_labels, true_labels)
specificty = torchmetrics.Specificity("binary").to(self.device)(predicted_labels, true_labels)
table = [
["0.5000", f"{acc.item():.4f}", f"{precision.item():.4f}", f"{recall.item():.4f}", f"{f1Score.item():.4f}", f"{auc.item():.4f}", f"{specificty.item():.4f}", f"*{self.current_epoch}*", "0"]
]
aucCurve = torchmetrics.ROC("binary").to(self.device)
fpr, tpr, thhols = aucCurve(predicted_labels, true_labels)
index = torch.argmax(tpr - fpr)
th1 = thhols[index].item()
accY = torchmetrics.Accuracy('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
precisionY = torchmetrics.Precision('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
recallY = torchmetrics.Recall('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
specifictyY = torchmetrics.Specificity("binary", threshold=th1).to(self.device)(predicted_labels, true_labels)
f1ScoreY = torchmetrics.F1Score('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
table.append([f"{th1:.4f}", f"{accY.item():.4f}", f"{precisionY.item():.4f}", f"{recallY.item():.4f}", f"{f1ScoreY.item():.4f}", f"{auc.item():.4f}", f"{specifictyY.item():.4f}", f"-{self.current_epoch}-", f"{(tpr - fpr)[index].item():.4f}"])
df = pd.DataFrame(table, columns=["Threshold", "Acc","Precision","Recall","F1","AUROC","Specificity","CheckPoint", "J"])
out_path = f"{self.logger.log_dir}/roc_points_epoch_{self.global_step}.csv"
df.to_csv(out_path, index=False)
class RARP_NVB_DINO_MultiTask_A6_Video(L.LightningModule):
def __init__(
self,
lr = 0.0001,
wd = 0.01,
L1 = None,
L2 = 0,
std = None,
mean = None,
head_type:int = 0, #None = 0, linear = 1, Attn. Pooling = 2, TCN = 3, Replace Head =4
chunks_loading:int = 50,
):
super().__init__()
self.lr = lr
self.wd = wd
self.chunks = chunks_loading
self.head_type = head_type
self.check_pt = True
self.base_model = van.van_b2(pretrained = True, num_classes = 0)
self.num_features_base_model = 512
#self.base_model_wrapper = ModuleWrapper(self.base_model)
self.lossFN = torch.nn.BCEWithLogitsLoss()
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
self.test_acc = torchmetrics.Accuracy('binary')
self.f1ScoreTest = torchmetrics.F1Score('binary')
self.test_probs = []
self.test_targets = []
match(self.head_type):
case 1:
self.head = TemporalConvNet(self.num_features_base_model, [128, 8, 1])
case _:
self.head = None
def _shared_video_step(self, batch:list[torch.Tensor], val_step:bool = False):
video, label = batch
video = video.float() #[B, T, C, H, W]
label = label.float() #[B]
pred_video = self(video, val_step) #[B, T, 1]
pred_video = pred_video.mean(dim=1) #[B, 1]
pred_video = pred_video.flatten() #[B] to match labels shape
predicted_labels = torch.sigmoid(pred_video)
loss = self.lossFN(pred_video, label)
return loss, label, predicted_labels
def forward(self, video:torch.Tensor, val_step:bool = True):
B, T, C, H, W = video.shape
chunk_T = self.chunks
pred_bt = []
def _fn(inp):
pred = self.base_model(inp)
return pred
for t0 in tqdm(range(0, T, chunk_T), desc=f"Video Analysis in {chunk_T} chunk", leave=False): # Loop for each chunk
t1 = min(T, t0 + chunk_T)
x = video[:, t0:t1].reshape(-1, C, H, W).contiguous(memory_format=torch.channels_last) # reshape from [B, chunk_T, C, H, W] to [B*chunk_T, C, H, W] and make the tensor GPU optimization
pred = torch_ckp.checkpoint(_fn, x) #froward to CNN and checkpoint grads
pred_bt.append(pred.view(B, t1-t0, -1)) #apped to output array and reshape to [B, chunk_T, C, H, W]
pred_video = torch.cat(pred_bt, dim=1) #concat all chunks -> [B, T, C, H, W]
pred_video = self.head(pred_video)
return pred_video
def on_after_backward(self):
norms = [p.grad.data.norm(2).item() for p in self.parameters() if p.grad is not None]
avg_layer_norm = sum(norms) / len(norms)
self.log("avg_grad_norm", avg_layer_norm)
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, False)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels = self._shared_video_step(batch, True)
self.test_probs.append(predicted_labels)
self.test_targets.append(true_labels)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def on_test_epoch_end(self):
predicted_labels = torch.cat(self.test_probs).to(self.device)
true_labels = torch.cat(self.test_targets).to(self.device).int()
acc = torchmetrics.Accuracy('binary').to(self.device)(predicted_labels, true_labels)
precision = torchmetrics.Precision('binary').to(self.device)(predicted_labels, true_labels)
recall = torchmetrics.Recall('binary').to(self.device)(predicted_labels, true_labels)
auc = torchmetrics.AUROC('binary').to(self.device)(predicted_labels, true_labels)
f1Score = torchmetrics.F1Score('binary').to(self.device)(predicted_labels, true_labels)
specificty = torchmetrics.Specificity("binary").to(self.device)(predicted_labels, true_labels)
table = [
["0.5000", f"{acc.item():.4f}", f"{precision.item():.4f}", f"{recall.item():.4f}", f"{f1Score.item():.4f}", f"{auc.item():.4f}", f"{specificty.item():.4f}", f"*{self.current_epoch}*", "0"]
]
aucCurve = torchmetrics.ROC("binary").to(self.device)
fpr, tpr, thhols = aucCurve(predicted_labels, true_labels)
index = torch.argmax(tpr - fpr)
th1 = thhols[index].item()
accY = torchmetrics.Accuracy('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
precisionY = torchmetrics.Precision('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
recallY = torchmetrics.Recall('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
specifictyY = torchmetrics.Specificity("binary", threshold=th1).to(self.device)(predicted_labels, true_labels)
f1ScoreY = torchmetrics.F1Score('binary', threshold=th1).to(self.device)(predicted_labels, true_labels)
table.append([f"{th1:.4f}", f"{accY.item():.4f}", f"{precisionY.item():.4f}", f"{recallY.item():.4f}", f"{f1ScoreY.item():.4f}", f"{auc.item():.4f}", f"{specifictyY.item():.4f}", f"-{self.current_epoch}-", f"{(tpr - fpr)[index].item():.4f}"])
df = pd.DataFrame(table, columns=["Threshold", "Acc","Precision","Recall","F1","AUROC","Specificity","CheckPoint", "J"])
out_path = f"{self.logger.log_dir}/roc_points_epoch_{self.global_step}.csv"
df.to_csv(out_path, index=False)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=self.wd) #, weight_decay=self.Lambda_L2
return [optimizer]
class RARP_NVB_DINO_MultiTask_Unet(RARP_NVB_DINO_MultiTask):
def _encoder_hool_fn(self, module, input, output):
self.feature_maps.append(output)
def _register_encoder_hooks(self):
for layer in self.list_blocks:
self.hooks.append(layer.register_forward_hook(self._encoder_hool_fn))
def __init__(
self,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher:float = 0.9995,
lr:float = 1e-4,
L1:float = None,
L2:float = 0,
std: float = None,
mean: float = None,
SoftAdptAlgo:int = 0,
SoftAdptBeta:float = 0.1
):
super().__init__(TypeLoss, momentum_teacher, lr, L1, L2, std, mean, SoftAdptAlgo, SoftAdptBeta)
self.hooks = []
self.feature_maps = []
self.list_blocks = [
self.student.backbone.block1[-1],
self.student.backbone.block2[-1],
self.student.backbone.block3[-1],
self.student.backbone.block4[-2],
]
self._register_encoder_hooks()
self.decoder = DecoderUnet(1024)
def forward(self, data, val_step:bool = True):
self.feature_maps = []
if val_step:
data = data.float()
dataTeacher, dataStudent = data, data
else:
data = [d.float() for d in data]
dataTeacher, dataStudent = data[1:3], data
TeacherDino = self.teacher_Features(dataTeacher)
Student = self.student(dataStudent)
if not val_step:
NumChunks = len(dataStudent)
S_GlogalViews = self.last_conv_output_S.chunk(NumChunks)[1:3]
self.last_conv_output_S = torch.cat(S_GlogalViews, dim=0)
cat_features = torch.cat((self.last_conv_output_S, self.last_conv_output_T), dim=1)
if not val_step:
NumChunks = len(dataStudent)
temp = []
for f_maps in self.feature_maps:
temp.append(torch.cat(f_maps.chunk(NumChunks)[1:3], dim=0))
self.feature_maps = temp
del temp
self.feature_maps.append(cat_features)
reconstructed_image = self.decoder(self.feature_maps)
Cont_Net = torch.nn.functional.adaptive_avg_pool2d(cat_features, (1,1)).flatten(1)
pred = self.clasiffier(Cont_Net)
return pred, (Student, TeacherDino), reconstructed_image
class RARP_NVB_DINO_MultiTask_MultiLabel(RARP_NVB_DINO_MultiTask):
def __init__(
self,
TypeLoss=TypeLossFunction.BCEWithLogits,
momentum_teacher = 0.9995,
lr = 0.0001,
L1 = None,
L2 = 0,
std = None,
mean = None,
SoftAdptAlgo = 0,
SoftAdptBeta = 0.1,
Num_Lables = 2
):
super().__init__(TypeLoss, momentum_teacher, lr, L1, L2, std, mean, SoftAdptAlgo, SoftAdptBeta)
self.lossFN = torch.nn.BCEWithLogitsLoss()
self.clasiffier = torch.nn.Sequential(
torch.nn.Linear(1024, 128),
torch.nn.Dropout(0.4),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8),
torch.nn.Dropout(0.2),
torch.nn.SiLU(True),
torch.nn.Linear(8, Num_Lables)
)
self.train_acc = torchmetrics.Accuracy('multilabel', num_labels=Num_Lables)
self.val_acc = torchmetrics.Accuracy('multilabel', num_labels=Num_Lables)
self.test_acc = torchmetrics.Accuracy('multilabel', num_labels=Num_Lables)
self.f1ScoreTest = torchmetrics.F1Score('multilabel', num_labels=Num_Lables)
class RARP_NVB_DINO_MultiTask_v2(RARP_NVB_DINO_MultiTask):
def __init__(
self,
TypeLoss=TypeLossFunction.CrossEntropy,
momentum_teacher = 0.9995,
lr = 0.0001,
L1 = None,
L2 = 0,
std = None,
mean = None
):
super().__init__(TypeLoss, momentum_teacher, lr, L1, L2, std, mean)
self.train_acc = torchmetrics.Accuracy('multiclass', num_classes=4)
self.val_acc = torchmetrics.Accuracy('multiclass', num_classes=4)
self.test_acc = torchmetrics.Accuracy('multiclass', num_classes=4)
self.f1ScoreTest = torchmetrics.F1Score('multiclass', num_classes=4)
self.lossFN = torch.nn.CrossEntropyLoss()
self.softAdapt = LossWeightedSoftAdapt(0.1) #NormalizedSoftAdapt(0.1)
self.clasiffier = torch.nn.Sequential(
torch.nn.Linear(1024, 512),
torch.nn.Dropout(0.2),
torch.nn.SiLU(True),
torch.nn.Linear(512, 256),
#torch.nn.Dropout(0.4),
torch.nn.SiLU(True),
torch.nn.Linear(256, 128),
#torch.nn.Dropout(0.4),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8),
#torch.nn.Dropout(0.2),
torch.nn.SiLU(True),
torch.nn.Linear(8, 4)
)
def _shared_step(self, batch, val_step:bool = False):
img, label = batch
prediction, features, new_image = self(img, val_step)
StudentF, TeacherF = features
predicted_labels = torch.softmax(prediction,dim=1)
orignalImg = torch.cat([img[0].float() for _ in range(len(TeacherF))], dim=0) if not val_step else img.float()
label = torch.cat([label for _ in range(len(TeacherF))], dim=0) if not val_step else label
#DINO Loss
loss_Dino = self.lossFN_DINO(StudentF, TeacherF) if not val_step else torch.tensor(0, device=label.device, dtype=torch.float32)
#Clasificator
loss_HL = self.lossFN(predicted_labels, label)
#Reconstruction
loss_img = self.ReconstructionLoss(new_image, orignalImg)
loss_img = loss_img.float()
if not val_step:
if self.Lambda_L1 is not None:
loss_l1 = 0
for params in self.clasiffier.parameters(): # aqui
loss_l1 += torch.sum(torch.abs(params))
loss_HL += self.Lambda_L1 * loss_l1
if self.Lambda_L2 > 0:
l2_reg = 0.0
for param in self.clasiffier.parameters():
l2_reg += torch.norm(param, 2) ** 2
loss_HL += self.Lambda_L2 * l2_reg
self.loss_history["loss_DINO"].append(loss_Dino.item())
self.loss_history["loss_Reconstruction"].append(loss_img.item())
self.loss_history["loss_Binary"].append(loss_HL.item())
loss = self.weights[0] * loss_Dino + self.weights[1] * loss_img + self.weights[2] * loss_HL
return loss, label, predicted_labels, (self.weights[0] * loss_Dino, self.weights[2] * loss_HL, self.weights[1] * loss_img, new_image)
class RARP_NVB_RN50_VAN_V2 (RARP_NVB_Model):
# Define a hook function to capture the output
def _hook_fn(self, module, input, output):
self.last_conv_output = output
def __init__(
self,
PseudoEstimator:str = None,
threshold:float = 0.5,
InitWeight=None,
TypeLoss=TypeLossFunction.CrossEntropy,
PseudoLables:bool=True,
HParameter = {},
std: float = None,
mean: float = None,
**kwargs
) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.std_IMG = torch.tensor(std).view(3, 1, 1) if std is not None else None
self.mean_IMG = torch.tensor(mean).view(3, 1, 1) if mean is not None else None
#self.RARP_RestNet50 = torch.hub.load('facebookresearch/dino:main', 'dino_resnet50')
self.RARP_RestNet50 = RARP_NVB_ResNet50.load_from_checkpoint(PseudoEstimator) if PseudoEstimator is not None else RARP_NVB_ResNet50()
self.RARP_RestNet50.model.fc = torch.nn.Identity()
for parms in self.RARP_RestNet50.model.parameters():
parms.requires_grad = False
self.RARP_RestNet50 = torch.nn.Sequential(*list(self.RARP_RestNet50.model.children())[:-2])
self.decoder = Decoder()
self.FeaturesLoss = FeatureAlignmentLoss()
self.ReconstructionLoss = torch.nn.MSELoss()
match (TypeError):
case TypeLossFunction.ContrastiveLoss:
self.lossFN = ContrastiveLoss()
case _:
pass
self.model = van.van_b2(pretrained = True, num_classes = 1)
self.Lambda_L1 = getNVL(HParameter, "L1", 1.31E-04)
self.lr = getNVL(HParameter, "lr", 1.0E-4)
print(f"lr={self.lr}, L1={self.Lambda_L1}")
# Initialize a variable to store the output
self.last_conv_output = None
self.model.block4[-1].mlp.dwconv.register_forward_hook(self._hook_fn)
def forward(self, data):
dataTeacher, dataStudent, _ = data
dataStudent = dataStudent.float()
dataTeacher = dataTeacher.float()
feature_Teacher = self.RARP_RestNet50 (dataTeacher)
#feature_Student = self.model.forward_features(dataStudent)
pred = self.model(dataStudent)
feature_Student = self.last_conv_output
#feature_Teacher = torch.nn.functional.adaptive_avg_pool1d(feature_Teacher, feature_Student.size(-1)) #Re-size output vector to macth VAN
cat_features = (feature_Student + feature_Teacher) / 2
reconstructed_image = self.decoder(cat_features)
feature_Student = torch.nn.functional.adaptive_avg_pool2d(feature_Student, (1, 1)).flatten(1)
feature_Teacher = torch.nn.functional.adaptive_avg_pool2d(feature_Teacher, (1, 1)).flatten(1)
return pred, (feature_Student, feature_Teacher), reconstructed_image
def _shared_step(self, batch):
img, label = batch
label = label.float()
orignalImg = img[2].float()
prediction, features, new_image = self(img)
prediction = prediction.flatten()
predicted_labels = torch.sigmoid(prediction)
#cosine similarity
loss_cosine = self.FeaturesLoss(features[0], features[1])
#Clasificator
loss_hl = self.lossFN(prediction, label)
#Reconstruction
loss_img = self.ReconstructionLoss(new_image, orignalImg)
loss_img = loss_img.float()
loss = loss_cosine + loss_hl + loss_img
if self.Lambda_L1 is not None:
loss_l1 = 0
for params in self.model.parameters():
loss_l1 += torch.sum(torch.abs(params))
loss += self.Lambda_L1 * loss_l1
return loss, label, predicted_labels, (loss_cosine, loss_hl, loss_img, new_image)
def _denormalize(self, tensor:torch.Tensor):
# Move mean and std to the same device as the input tensor
mean = self.mean_IMG.to(tensor.device)
std = self.std_IMG.to(tensor.device)
return tensor * std + mean
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
self.log("train_loss_img", losses[2], on_epoch=True, on_step=False)
self.log("train_loss_cosSim", losses[0], on_epoch=True, on_step=False)
self.log("train_loss_GT", losses[1], on_epoch=True, on_step=False)
if batch_idx % 50 == 0 and self.mean_IMG is not None and self.std_IMG is not None:
imgReconstruction = torch.clip(self._denormalize(losses[3]) / 255, 0, 1)
imgReconstruction = imgReconstruction[:, [2, 1, 0], :, :]
grid = torchvision.utils.make_grid(imgReconstruction)
self.logger.experiment.add_image('reconstructed_images', grid, self.global_step)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_loss_img", losses[2], on_epoch=True, on_step=False)
self.log("val_loss_cosSim", losses[0], on_epoch=True, on_step=False)
self.log("val_loss_GT", losses[1], on_epoch=True, on_step=False)
def test_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
if self.mean_IMG is not None and self.std_IMG is not None:
imgReconstruction = torch.clip(self._denormalize(losses[3]) / 255, 0, 1)
imgReconstruction = imgReconstruction[:, [2, 1, 0], :, :]
grid = torchvision.utils.make_grid(imgReconstruction)
self.logger.experiment.add_image('reconstructed_images_test', grid, self.global_step)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.Lambda_L2)
return optimizer
class RARP_NVB_ResNet50_VAN(RARP_NVB_Model):
def __init__(
self,
PseudoEstimator:str = None,
threshold:float = 0.5,
InitWeight=None,
TypeLoss=TypeLossFunction.CrossEntropy,
PseudoLables:bool=True,
HParameter = {},
**kwargs
) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.RARP_RestNet50 = RARP_NVB_ResNet50.load_from_checkpoint(PseudoEstimator) if PseudoEstimator is not None else RARP_NVB_ResNet50()
self.threshold = threshold
self.PseudoLables = PseudoLables
match (TypeError):
case TypeLossFunction.ContrastiveLoss:
self.lossFN = ContrastiveLoss()
case _:
pass
for parms in self.RARP_RestNet50.model.parameters():
parms.requires_grad = False
self.model = van.van_b2(pretrained = True, num_classes = 1)
self.Lambda_L1 = getNVL(HParameter, "L1", 1.31E-04)
self.lr = getNVL(HParameter, "lr", 1.0E-4)
self.W_Apha = getNVL(HParameter, "Alpha", 0.60)
self.W_Beta = 1 - getNVL(HParameter, "Alpha", 0.60)
self.thao_KD = getNVL(HParameter, "Thao", 5)
print(f"A={self.W_Apha}; B={self.W_Beta}, T={self.thao_KD}, lr={self.lr}, L1={self.Lambda_L1}")
#self.Lambda_L1 = None
#self.scheduler = True
#self.lr = 1.74E-2
def forward(self, data):
#dataTeacher, dataStudent = data
if isinstance(data, tuple):
dataTeacher, dataStudent = data
elif isinstance(data, torch.Tensor):
dataTeacher, dataStudent = (data, data)
RN50Pred = torch.sigmoid(self.RARP_RestNet50(dataTeacher).flatten())
PseudoLabels = (RN50Pred > self.threshold) * 1.0
dataStudent = dataStudent.float()
pred = self.model(dataStudent)
return pred, PseudoLabels, RN50Pred
def _shared_step(self, batch):
img, label = batch
if self.InitWeight is not None:
self.lossFN.weight = self.InitWeight[label]
label = label.float()
prediction, PseudoLabels, predictionRN50 = self(img) #.flatten()
prediction = prediction.flatten()
predicted_labels = torch.sigmoid(prediction)
#
# L[B] => BCEWithLogits
# L[C] => ContrastiveLoss
#Loss L[1] = L[B_HL] + L[B_PL]
#Loss L[2] = L[C_HL]
#Loss L[3] = L[C_HL] + L[C_PL]
#Loss L[4] = L[C_y_hat]
#Loss L[5] = L[KLD_PL] + L[B_PL]
#Loss L[6] = L[MSE_PL] + L[B_PL]
#Loss L[7] = L[B_HL] + L[BCE_PL]
#Loss L[8] = FocalLoss*L[JSD]
#Nuveo
#thao_KD = 5.0
#w_alpha, w_beta = (0.60, 0.40)
#L[7]:
softTeacher = torch.sigmoid(predictionRN50/self.thao_KD)
softStudent = torch.sigmoid(prediction/self.thao_KD)
loss_sl = torch.nn.functional.binary_cross_entropy(softStudent, softTeacher)
loss_hl = self.lossFN(prediction, label)
loss = self.W_Apha * loss_hl + self.W_Beta * loss_sl #/ (self.thao_KD ** 2)
#loss = w_alpha * self.lossFN(prediction, label) + w_beta * self.lossFN(prediction, PseudoLabels) #L[1]
#loss = w_alpha * (torch.nn.KLDivLoss()(softStudent, softTeacher) * (thao_KD ** 2)) + w_beta * self.lossFN(prediction, PseudoLabels) #L[5]
#loss = w_alpha * torch.nn.functional.mse_loss(softStudent, softTeacher) + w_beta * self.lossFN(prediction, PseudoLabels) #L[6]
#loss = self.lossFN(predictionRN50, predicted_labels, label) #L[2]
#loss = self.lossFN(predictionRN50, predicted_labels, label) + self.lossFN(predictionRN50, predicted_labels, PseudoLabels) #L[3]
#y_hat = (PseudoLabels != label) * 1
#loss = self.lossFN(predictionRN50, predicted_labels, y_hat) #L[4]
if self.Lambda_L1 is not None:
loss_l1 = 0
for params in self.model.parameters():
loss_l1 += torch.sum(torch.abs(params))
loss += self.Lambda_L1 * loss_l1
return loss, (PseudoLabels if self.PseudoLables else label), predicted_labels, (loss_hl, loss_sl)
def training_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch)
self.log("train_loss", loss, on_epoch=True)
self.train_acc.update(predicted_labels, true_labels)
self.log("train_acc", self.train_acc, on_epoch=True, on_step=False)
self.log("train_lh", losses[0], on_step=False, on_epoch=True)
self.log("train_ld", losses[1], on_step=False, on_epoch=True)
return loss
def validation_step(self, batch, batch_idx):
loss, true_labels, predicted_labels, losses = self._shared_step(batch)
self.log("val_loss", loss, on_epoch=True, on_step=False, prog_bar=True)
self.val_acc.update(predicted_labels, true_labels)
self.log("val_acc", self.val_acc, on_epoch=True, on_step=False, prog_bar=True)
self.log("val_lh", losses[0], on_step=False, on_epoch=True)
self.log("val_ld", losses[1], on_step=False, on_epoch=True)
def test_step(self, batch, batch_idx):
_, true_labels, predicted_labels, _ = self._shared_step(batch)
self.test_acc.update(predicted_labels, true_labels)
self.f1ScoreTest.update(predicted_labels, true_labels)
self.log("test_acc", self.test_acc, on_epoch=True, on_step=False)
self.log("test_f1", self.f1ScoreTest, on_epoch=True, on_step=False)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.Lambda_L2)
if self.scheduler:
scheduler = CosineAnnealingLR(
optimizer,
max_epochs=150,
warmup_epochs=8,
warmup_start_lr=0.001,
eta_min=0.0001
)
return {
"optimizer": optimizer,
"lr_scheduler": {
"scheduler": scheduler,
#"monitor": "val_loss",
}
}
return optimizer
class RARP_NVB_SSL_RestNet50_Deep(RARP_NVB_ResNet50_VAN):
def __init__(
self,
PseudoEstimator: str = None,
threshold: float = 0.5,
InitWeight=None,
TypeLoss=TypeLossFunction.CrossEntropy,
PseudoLables: bool = True,
**kwargs
) -> None:
super().__init__(None, threshold, InitWeight, TypeLoss, PseudoLables, **kwargs)
self.RARP_RestNet50 = RARP_NVB_ResNet50_Deep.load_from_checkpoint(PseudoEstimator, strict=False) if PseudoEstimator is not None else RARP_NVB_ResNet50_Deep()
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
#self.model = torch.hub.load('facebookresearch/dino:main', 'dino_resnet50')
tempFC_ft = 2048
self.model.fc = torch.nn.Sequential(
torch.nn.Dropout(0.2),
torch.nn.Linear(in_features=tempFC_ft, out_features=128),
torch.nn.SiLU(True),
torch.nn.Linear(128, 8),
torch.nn.SiLU(True),
torch.nn.Linear(8, 1)
)
for parms in self.RARP_RestNet50.model.parameters():
parms.requires_grad = False
class RARP_NVB_MobileNetV2(RARP_NVB_Model):#class RARP_NVB_MobileNetV2(RARP_NVB_Model_BCEWithLogitsLoss):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.mobilenet_v2(weights=torchvision.models.MobileNet_V2_Weights.DEFAULT)
tempFC_ft = self.model.classifier[1].in_features
self.model.classifier[1] = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_EfficientNetV2(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy, **kwargs) -> None:
super().__init__(InitWeight, TypeLoss, **kwargs)
self.model = torchvision.models.efficientnet_v2_s(weights=torchvision.models.EfficientNet_V2_S_Weights.DEFAULT)
tempFC_ft = self.model.classifier[1].in_features
self.model.classifier[1] = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_Vit_b_16(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy) -> None:
super().__init__(InitWeight, TypeLoss)
self.model = torchvision.models.vit_b_16(weights=torchvision.models.ViT_B_16_Weights.DEFAULT)
tempFC_ft = self.model.heads[0].in_features
self.model.heads[0] = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
class RARP_NVB_DenseNet169(RARP_NVB_Model):
def __init__(self, InitWeight=None, TypeLoss=TypeLossFunction.CrossEntropy) -> None:
super().__init__(InitWeight, TypeLoss)
self.model = torchvision.models.densenet169(weights=torchvision.models.DenseNet169_Weights.DEFAULT)
inFeatures = self.model.classifier.in_features
self.model.classifier = torch.nn.Linear(in_features=inFeatures, out_features=1)
class RARP_NVB_RestNet50_old(L.LightningModule):
def __init__(self, InitialWeigth = 1) -> None:
super().__init__()
self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
tempFC_ft = self.model.fc.in_features
self.model.fc = torch.nn.Linear(in_features=tempFC_ft, out_features=1)
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
self.lossFN = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([InitialWeigth]))
self.train_acc = torchmetrics.Accuracy('binary')
self.val_acc = torchmetrics.Accuracy('binary')
def forward(self, data):
data = data.float()
pred = self.model(data)
return pred
def training_step(self, batch, batch_idx):
img, label = batch
label = label.float()
prediction = self(img)[:,0]
loss = self.lossFN(prediction, label)
self.log("Train Loss", loss)
self.log("Step Train Acc", self.train_acc(torch.sigmoid(prediction), label.int()))
return loss
def on_train_epoch_end(self):
self.log("Train Acc", self.train_acc.compute())
def validation_step(self, batch, batch_idx):
img, label = batch
label = label.float()
prediction = self(img)[:,0]
loss = self.lossFN(prediction, label)
self.log("Val Loss", loss)
self.log("Step Val Acc", self.val_acc(torch.sigmoid(prediction), label.int())) #train_acc
return loss
def on_validation_epoch_end(self):
self.log("Val Acc", self.val_acc.compute())
def configure_optimizers(self):
return [self.optimizer]
