|
import torchmetrics |
|
from . import config |
|
|
|
from typing import Tuple, Dict, List, Any |
|
|
|
import numpy as np |
|
import torch |
|
import torchvision |
|
import torch.nn as nn |
|
import pytorch_lightning as ptl |
|
|
|
|
|
class DeepFontBaseline(nn.Module): |
|
def __init__(self) -> None: |
|
super().__init__() |
|
self.model = nn.Sequential( |
|
nn.Conv2d(3, 64, 11, 2), |
|
nn.BatchNorm2d(64), |
|
nn.ReLU(), |
|
nn.MaxPool2d(2, 2), |
|
nn.Conv2d(64, 128, 3, 1, 1), |
|
nn.BatchNorm2d(128), |
|
nn.ReLU(), |
|
nn.MaxPool2d(2, 2), |
|
nn.Conv2d(128, 256, 3, 1, 1), |
|
nn.ReLU(), |
|
nn.Conv2d(256, 256, 3, 1, 1), |
|
nn.ReLU(), |
|
nn.Conv2d(256, 256, 3, 1, 1), |
|
nn.ReLU(), |
|
|
|
nn.Flatten(), |
|
nn.Linear(256 * 12 * 12, 4096), |
|
nn.ReLU(), |
|
nn.Linear(4096, 4096), |
|
nn.ReLU(), |
|
nn.Linear(4096, config.FONT_COUNT), |
|
) |
|
|
|
def forward(self, X): |
|
return self.model(X) |
|
|
|
|
|
class ResNet18Regressor(nn.Module): |
|
def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False): |
|
super().__init__() |
|
weights = torchvision.models.ResNet18_Weights.DEFAULT if pretrained else None |
|
self.model = torchvision.models.resnet18(weights=weights) |
|
self.model.fc = nn.Linear(512, config.FONT_COUNT + 12) |
|
self.regression_use_tanh = regression_use_tanh |
|
|
|
def forward(self, X): |
|
X = self.model(X) |
|
|
|
if not self.regression_use_tanh: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].sigmoid() |
|
else: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].tanh() |
|
return X |
|
|
|
|
|
class ResNet34Regressor(nn.Module): |
|
def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False): |
|
super().__init__() |
|
weights = torchvision.models.ResNet34_Weights.DEFAULT if pretrained else None |
|
self.model = torchvision.models.resnet34(weights=weights) |
|
self.model.fc = nn.Linear(512, config.FONT_COUNT + 12) |
|
self.regression_use_tanh = regression_use_tanh |
|
|
|
def forward(self, X): |
|
X = self.model(X) |
|
|
|
if not self.regression_use_tanh: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].sigmoid() |
|
else: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].tanh() |
|
return X |
|
|
|
|
|
class ResNet50Regressor(nn.Module): |
|
def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False): |
|
super().__init__() |
|
weights = torchvision.models.ResNet50_Weights.DEFAULT if pretrained else None |
|
self.model = torchvision.models.resnet50(weights=weights) |
|
self.model.fc = nn.Linear(2048, config.FONT_COUNT + 12) |
|
self.regression_use_tanh = regression_use_tanh |
|
|
|
def forward(self, X): |
|
X = self.model(X) |
|
|
|
if not self.regression_use_tanh: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].sigmoid() |
|
else: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].tanh() |
|
return X |
|
|
|
|
|
class ResNet101Regressor(nn.Module): |
|
def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False): |
|
super().__init__() |
|
weights = torchvision.models.ResNet101_Weights.DEFAULT if pretrained else None |
|
self.model = torchvision.models.resnet101(weights=weights) |
|
self.model.fc = nn.Linear(2048, config.FONT_COUNT + 12) |
|
self.regression_use_tanh = regression_use_tanh |
|
|
|
def forward(self, X): |
|
X = self.model(X) |
|
|
|
if not self.regression_use_tanh: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].sigmoid() |
|
else: |
|
X[..., config.FONT_COUNT + 2 :] = X[..., config.FONT_COUNT + 2 :].tanh() |
|
return X |
|
|
|
|
|
class FontDetectorLoss(nn.Module): |
|
def __init__( |
|
self, lambda_font, lambda_direction, lambda_regression, font_classification_only |
|
): |
|
super().__init__() |
|
self.category_loss = nn.CrossEntropyLoss() |
|
self.regression_loss = nn.MSELoss() |
|
self.lambda_font = lambda_font |
|
self.lambda_direction = lambda_direction |
|
self.lambda_regression = lambda_regression |
|
self.font_classfiication_only = font_classification_only |
|
|
|
def forward(self, y_hat, y): |
|
font_cat = self.category_loss(y_hat[..., : config.FONT_COUNT], y[..., 0].long()) |
|
if self.font_classfiication_only: |
|
return self.lambda_font * font_cat |
|
direction_cat = self.category_loss( |
|
y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1].long() |
|
) |
|
regression = self.regression_loss( |
|
y_hat[..., config.FONT_COUNT + 2 :], y[..., 2:] |
|
) |
|
return ( |
|
self.lambda_font * font_cat |
|
+ self.lambda_direction * direction_cat |
|
+ self.lambda_regression * regression |
|
) |
|
|
|
|
|
class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler): |
|
def __init__(self, optimizer, warmup, max_iters): |
|
self.warmup = warmup |
|
self.max_num_iters = max_iters |
|
super().__init__(optimizer) |
|
|
|
def get_lr(self): |
|
lr_factor = self.get_lr_factor(epoch=self.last_epoch) |
|
return [base_lr * lr_factor for base_lr in self.base_lrs] |
|
|
|
def get_lr_factor(self, epoch): |
|
lr_factor = 0.5 * (1 + np.cos(np.pi * epoch / self.max_num_iters)) |
|
if epoch <= self.warmup: |
|
lr_factor *= epoch * 1.0 / self.warmup |
|
return lr_factor |
|
|
|
|
|
class FontDetector(ptl.LightningModule): |
|
def __init__( |
|
self, |
|
model: nn.Module, |
|
lambda_font: float, |
|
lambda_direction: float, |
|
lambda_regression: float, |
|
font_classification_only: bool, |
|
lr: float, |
|
betas: Tuple[float, float], |
|
num_warmup_iters: int, |
|
num_iters: int, |
|
num_epochs: int, |
|
): |
|
super().__init__() |
|
self.model = model |
|
self.loss = FontDetectorLoss( |
|
lambda_font, lambda_direction, lambda_regression, font_classification_only |
|
) |
|
self.font_accur_train = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=config.FONT_COUNT |
|
) |
|
self.font_accur_val = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=config.FONT_COUNT |
|
) |
|
self.font_accur_test = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=config.FONT_COUNT |
|
) |
|
if not font_classification_only: |
|
self.direction_accur_train = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=2 |
|
) |
|
self.direction_accur_val = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=2 |
|
) |
|
self.direction_accur_test = torchmetrics.Accuracy( |
|
task="multiclass", num_classes=2 |
|
) |
|
self.lr = lr |
|
self.betas = betas |
|
self.num_warmup_iters = num_warmup_iters |
|
self.num_iters = num_iters |
|
self.num_epochs = num_epochs |
|
self.load_epoch = 0 |
|
self.font_classification_only = font_classification_only |
|
|
|
def forward(self, x): |
|
return self.model(x) |
|
|
|
def training_step( |
|
self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int |
|
) -> Dict[str, Any]: |
|
X, y = batch |
|
y_hat = self.forward(X) |
|
loss = self.loss(y_hat, y) |
|
self.log("train_loss", loss, prog_bar=True, sync_dist=True) |
|
|
|
self.log( |
|
"train_font_accur", |
|
self.font_accur_train(y_hat[..., : config.FONT_COUNT], y[..., 0]), |
|
sync_dist=True, |
|
) |
|
if not self.font_classification_only: |
|
self.log( |
|
"train_direction_accur", |
|
self.direction_accur_train( |
|
y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1] |
|
), |
|
sync_dist=True, |
|
) |
|
return {"loss": loss} |
|
|
|
def on_train_epoch_end(self) -> None: |
|
self.log("train_font_accur", self.font_accur_train.compute(), sync_dist=True) |
|
self.font_accur_train.reset() |
|
if not self.font_classification_only: |
|
self.log( |
|
"train_direction_accur", |
|
self.direction_accur_train.compute(), |
|
sync_dist=True, |
|
) |
|
self.direction_accur_train.reset() |
|
|
|
def validation_step( |
|
self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int |
|
) -> Dict[str, Any]: |
|
X, y = batch |
|
y_hat = self.forward(X) |
|
loss = self.loss(y_hat, y) |
|
self.log("val_loss", loss, prog_bar=True, sync_dist=True) |
|
self.font_accur_val.update(y_hat[..., : config.FONT_COUNT], y[..., 0]) |
|
if not self.font_classification_only: |
|
self.direction_accur_val.update( |
|
y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1] |
|
) |
|
return {"loss": loss} |
|
|
|
def on_validation_epoch_end(self): |
|
self.log("val_font_accur", self.font_accur_val.compute(), sync_dist=True) |
|
self.font_accur_val.reset() |
|
if not self.font_classification_only: |
|
self.log( |
|
"val_direction_accur", |
|
self.direction_accur_val.compute(), |
|
sync_dist=True, |
|
) |
|
self.direction_accur_val.reset() |
|
|
|
def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int): |
|
X, y = batch |
|
y_hat = self.forward(X) |
|
loss = self.loss(y_hat, y) |
|
self.log("test_loss", loss, prog_bar=True, sync_dist=True) |
|
self.font_accur_test.update(y_hat[..., : config.FONT_COUNT], y[..., 0]) |
|
if not self.font_classification_only: |
|
self.direction_accur_test.update( |
|
y_hat[..., config.FONT_COUNT : config.FONT_COUNT + 2], y[..., 1] |
|
) |
|
return {"loss": loss} |
|
|
|
def on_test_epoch_end(self) -> None: |
|
self.log("test_font_accur", self.font_accur_test.compute(), sync_dist=True) |
|
self.font_accur_test.reset() |
|
if not self.font_classification_only: |
|
self.log( |
|
"test_direction_accur", |
|
self.direction_accur_test.compute(), |
|
sync_dist=True, |
|
) |
|
self.direction_accur_test.reset() |
|
|
|
def configure_optimizers(self): |
|
optimizer = torch.optim.Adam( |
|
self.model.parameters(), lr=self.lr, betas=self.betas |
|
) |
|
self.scheduler = CosineWarmupScheduler( |
|
optimizer, self.num_warmup_iters, self.num_iters |
|
) |
|
print("Load epoch:", self.load_epoch) |
|
for _ in range(self.num_iters * self.load_epoch // self.num_epochs): |
|
self.scheduler.step() |
|
print("Current learning rate set to:", self.scheduler.get_last_lr()) |
|
return optimizer |
|
|
|
def optimizer_step( |
|
self, |
|
epoch: int, |
|
batch_idx: int, |
|
optimizer, |
|
optimizer_idx: int = 0, |
|
*args, |
|
**kwargs |
|
): |
|
super().optimizer_step( |
|
epoch, batch_idx, optimizer, optimizer_idx, *args, **kwargs |
|
) |
|
self.log("lr", self.scheduler.get_last_lr()[0]) |
|
self.scheduler.step() |
|
|
|
def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None: |
|
self.load_epoch = checkpoint["epoch"] |
|
|