Обнавлена структура проекта
This commit is contained in:
Даниил Грабарь
@@ -0,0 +1,2 @@
|
||||
EXPERIMENT_NAME=
|
||||
TRACKING_URI=
|
||||
+7
-1
@@ -1,4 +1,10 @@
|
||||
.vscode/
|
||||
__pycache__/
|
||||
|
||||
models/
|
||||
Datasets/
|
||||
test_images/
|
||||
train_logs/
|
||||
train_logs/
|
||||
lightning_logs/
|
||||
|
||||
.env
|
||||
@@ -1,488 +0,0 @@
|
||||
# %% Импорт библиотек
|
||||
|
||||
import time
|
||||
import torch
|
||||
import kagglehub
|
||||
import pandas as pd
|
||||
import torch.nn as nn
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from PIL import Image
|
||||
from torch import optim
|
||||
from pathlib import Path
|
||||
from torchmetrics import MetricCollection
|
||||
from typing import Optional, Callable, Any
|
||||
from torchvision import datasets, transforms
|
||||
from torchmetrics.classification import Accuracy
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
# Логирование данных
|
||||
from pytorch_lightning.loggers import MLFlowLogger
|
||||
# автоматическое отслеживание lr
|
||||
from pytorch_lightning.callbacks import LearningRateMonitor
|
||||
# ранние остановки
|
||||
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
|
||||
# сохранение весов
|
||||
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
|
||||
# %% Параметры
|
||||
|
||||
dataset_root = "/media/daniel/developments/Datasets/ImageNet"
|
||||
batch_size = 128
|
||||
num_workers = 8
|
||||
num_epoch = 90
|
||||
learning_rate = 3e-4
|
||||
|
||||
# %% Трансформации
|
||||
|
||||
train_transform = transforms.Compose([
|
||||
transforms.RandomResizedCrop(224),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
|
||||
transforms.ToTensor(),
|
||||
|
||||
transforms.Normalize(
|
||||
mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225]
|
||||
)
|
||||
])
|
||||
|
||||
eval_transform = transforms.Compose([
|
||||
transforms.Resize(256),
|
||||
transforms.CenterCrop(224),
|
||||
|
||||
transforms.ToTensor(),
|
||||
|
||||
transforms.Normalize(
|
||||
mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225]
|
||||
)
|
||||
])
|
||||
|
||||
# %% Скачивание ImageNet
|
||||
|
||||
path = kagglehub.competition_download('imagenet-object-localization-challenge',
|
||||
output_dir=dataset_root,
|
||||
force_download=False)
|
||||
print("Path to competition files:", path)
|
||||
# %% Датасет
|
||||
|
||||
TransformType = Optional[Callable[..., Any]]
|
||||
|
||||
class ImageNetValDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
image_dir,
|
||||
labels_file,
|
||||
class_to_idx,
|
||||
classes,
|
||||
transform=None
|
||||
):
|
||||
self.image_dir = Path(image_dir)
|
||||
self.transform = transform
|
||||
|
||||
self.class_to_idx = class_to_idx
|
||||
self.classes = classes
|
||||
|
||||
df = pd.read_csv(labels_file)
|
||||
|
||||
self.samples = []
|
||||
|
||||
for _, row in df.iterrows():
|
||||
image_id = row["ImageId"]
|
||||
synset = row["PredictionString"].split()[0]
|
||||
label = self.class_to_idx[synset]
|
||||
image_path = self.image_dir / f"{image_id}.JPEG"
|
||||
self.samples.append(
|
||||
(image_path, label)
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_path, label = self.samples[idx]
|
||||
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
|
||||
if self.transform is not None:
|
||||
image = self.transform(image)
|
||||
|
||||
return image, label
|
||||
|
||||
|
||||
class ImageNetDataModule(pl.LightningDataModule):
|
||||
def __init__(self,
|
||||
dataset_root: Path,
|
||||
train_transform: TransformType,
|
||||
eval_transform: TransformType,
|
||||
batch_size: int = 64,
|
||||
num_workers:int = 8
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dataset_root = Path(dataset_root)
|
||||
|
||||
self.train_transform = train_transform
|
||||
self.eval_transform = eval_transform
|
||||
self.num_workers = num_workers
|
||||
self.batch_size = batch_size
|
||||
|
||||
self.imagenet_root = Path(dataset_root) / "ILSVRC"
|
||||
|
||||
self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
|
||||
self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
|
||||
|
||||
self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
|
||||
self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"
|
||||
|
||||
self.train_dataset = None
|
||||
self.val_dataset = None
|
||||
|
||||
self.class_to_idx = None
|
||||
self.classes = None
|
||||
|
||||
def setup(self, stage=None):
|
||||
synset_to_name = {}
|
||||
with open(self.mapping_file, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
synset, class_name = line.strip().split(" ", 1)
|
||||
synset_to_name[synset] = class_name
|
||||
|
||||
train_dataset = datasets.ImageFolder(root=self.train_dir,
|
||||
transform=self.train_transform)
|
||||
|
||||
self.class_to_idx = train_dataset.class_to_idx.copy()
|
||||
self.classes = [synset_to_name[synset]
|
||||
for synset in train_dataset.classes]
|
||||
train_dataset.classes = self.classes
|
||||
self.train_dataset = train_dataset
|
||||
|
||||
self.val_dataset = ImageNetValDataset(
|
||||
image_dir=self.val_dir,
|
||||
labels_file=self.val_labels_file,
|
||||
class_to_idx=self.class_to_idx,
|
||||
classes=self.classes,
|
||||
transform=self.eval_transform
|
||||
)
|
||||
|
||||
def train_dataloader(self):
|
||||
return DataLoader(
|
||||
dataset=self.train_dataset,
|
||||
batch_size=self.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=self.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=self.num_workers > 0,
|
||||
prefetch_factor=4 if self.num_workers > 0 else None
|
||||
)
|
||||
|
||||
def val_dataloader(self):
|
||||
return DataLoader(
|
||||
dataset=self.val_dataset,
|
||||
batch_size=self.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=self.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=self.num_workers > 0,
|
||||
prefetch_factor=4 if self.num_workers > 0 else None
|
||||
)
|
||||
|
||||
def test_dataloader(self):
|
||||
return DataLoader(
|
||||
dataset=self.val_dataset,
|
||||
batch_size=self.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=self.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=self.num_workers > 0
|
||||
)
|
||||
|
||||
# %% AlexNet
|
||||
|
||||
class AlexNetCustom(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels: int = 3,
|
||||
out_classes: int=1_000):
|
||||
super().__init__()
|
||||
|
||||
self._out_classes = out_classes
|
||||
|
||||
self.feature = nn.Sequential(
|
||||
nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2),
|
||||
|
||||
nn.Conv2d(96, 256, kernel_size=5, padding=2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2),
|
||||
|
||||
nn.Conv2d(256, 384, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(384, 384, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(384, 256, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2)
|
||||
)
|
||||
|
||||
# Добавлено в современной реализации
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Dropout(p=0.5),
|
||||
nn.Linear(256*6*6, 4096),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Dropout(p=0.5),
|
||||
nn.Linear(4096, 4096),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Linear(4096, self._out_classes)
|
||||
)
|
||||
|
||||
def forward(self, image):
|
||||
x = self.feature(image)
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.mlp(x)
|
||||
|
||||
return x
|
||||
|
||||
# %% Классификатор обучатор
|
||||
|
||||
class ClassificationLightning(pl.LightningModule):
|
||||
def __init__(self,
|
||||
model: nn.Module,
|
||||
num_classes: int,
|
||||
optimizer_class,
|
||||
optimizer_kwargs: dict,
|
||||
scheduler_class=None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
scheduler_monitor: str = "val/loss"):
|
||||
super().__init__()
|
||||
|
||||
|
||||
scheduler_kwargs = scheduler_kwargs or {}
|
||||
|
||||
self.model = model
|
||||
self.optimizer_class = optimizer_class
|
||||
self.scheduler_class = scheduler_class
|
||||
|
||||
num_parameters = sum(
|
||||
p.numel()
|
||||
for p in model.parameters()
|
||||
)
|
||||
|
||||
num_trainable_parameters = sum(
|
||||
p.numel()
|
||||
for p in model.parameters()
|
||||
if p.requires_grad
|
||||
)
|
||||
|
||||
self.save_hyperparameters({
|
||||
"model": model.__class__.__name__,
|
||||
"num_classes": num_classes,
|
||||
|
||||
"num_parameters": num_parameters,
|
||||
"num_trainable_parameters": num_trainable_parameters,
|
||||
|
||||
"optimizer": optimizer_class.__name__,
|
||||
"optimizer_kwargs": optimizer_kwargs,
|
||||
|
||||
"scheduler": (
|
||||
scheduler_class.__name__
|
||||
if scheduler_class is not None
|
||||
else None
|
||||
),
|
||||
"scheduler_kwargs": scheduler_kwargs,
|
||||
"scheduler_monitor": scheduler_monitor
|
||||
})
|
||||
|
||||
|
||||
self.criterion = nn.CrossEntropyLoss()
|
||||
self.train_metrics = MetricCollection({
|
||||
"acc_top1": Accuracy(task="multiclass",
|
||||
num_classes=self.hparams.num_classes,
|
||||
top_k=1),
|
||||
"acc_top5": Accuracy(task="multiclass",
|
||||
num_classes=self.hparams.num_classes,
|
||||
top_k=5),
|
||||
})
|
||||
|
||||
self.val_metrics = self.train_metrics.clone()
|
||||
self.test_metrics = self.train_metrics.clone()
|
||||
|
||||
def forward(self, image):
|
||||
return self.model(image)
|
||||
|
||||
def _shared_step(self,
|
||||
sample,
|
||||
stage,
|
||||
metrics):
|
||||
x, y = sample
|
||||
|
||||
logits = self(x)
|
||||
loss = self.criterion(logits, y)
|
||||
|
||||
metrics.update(logits, y)
|
||||
|
||||
self.log(f"{stage}/loss",
|
||||
loss,
|
||||
prog_bar=True,
|
||||
on_epoch=True,
|
||||
on_step=False
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
def _log_metrics(self, stage, metrics):
|
||||
computed = metrics.compute()
|
||||
|
||||
for name, value in computed.items():
|
||||
self.log(f"{stage}/{name}", value, prog_bar=True)
|
||||
|
||||
metrics.reset()
|
||||
|
||||
def on_train_epoch_start(self):
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self._epoch_start_time = time.perf_counter()
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "train", self.train_metrics)
|
||||
|
||||
def on_train_epoch_end(self):
|
||||
self._log_metrics("train", self.train_metrics)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
epoch_time = time.perf_counter() - self._epoch_start_time
|
||||
lr = self.optimizers().param_groups[0]["lr"]
|
||||
|
||||
self.log("train/lr", lr, prog_bar=True)
|
||||
self.log("train/epoch_time", epoch_time)
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "val", self.val_metrics)
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
self._log_metrics("val", self.val_metrics)
|
||||
|
||||
def test_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "test", self.test_metrics)
|
||||
|
||||
def on_test_epoch_end(self):
|
||||
self._log_metrics("test", self.test_metrics)
|
||||
|
||||
def configure_optimizers(self):
|
||||
optimizer = self.optimizer_class(
|
||||
self.parameters(),
|
||||
**self.hparams.optimizer_kwargs
|
||||
)
|
||||
|
||||
if self.scheduler_class is None:
|
||||
return optimizer
|
||||
|
||||
scheduler = self.scheduler_class(
|
||||
optimizer,
|
||||
**self.hparams.scheduler_kwargs
|
||||
)
|
||||
|
||||
scheduler_config = {
|
||||
"scheduler": scheduler,
|
||||
"interval": "epoch",
|
||||
"frequency": 1
|
||||
}
|
||||
|
||||
if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
|
||||
scheduler_config["monitor"] = self.hparams.scheduler_monitor
|
||||
|
||||
return {
|
||||
"optimizer": optimizer,
|
||||
"lr_scheduler": scheduler_config
|
||||
}
|
||||
|
||||
# %% Обучение
|
||||
|
||||
pl.seed_everything(42, workers=True)
|
||||
torch.set_float32_matmul_precision("high")
|
||||
|
||||
imagenet_dm = ImageNetDataModule(
|
||||
dataset_root=dataset_root,
|
||||
train_transform=train_transform,
|
||||
eval_transform=eval_transform,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers
|
||||
)
|
||||
|
||||
|
||||
optimizer_Adam = optim.Adam
|
||||
optimizer_Adam_kwargs = {
|
||||
"lr": 3e-4,
|
||||
"weight_decay": 1e-4
|
||||
}
|
||||
|
||||
scheduler_ReduceLR = optim.lr_scheduler.ReduceLROnPlateau
|
||||
scheduler_ReduceLR_kwargs={
|
||||
"mode": "min",
|
||||
"factor": 0.1,
|
||||
"patience": 5
|
||||
}
|
||||
|
||||
callbacks = [
|
||||
ModelCheckpoint(
|
||||
dirpath="models",
|
||||
filename="epoch={epoch}",
|
||||
auto_insert_metric_name=False,
|
||||
save_last=True,
|
||||
save_top_k=1,
|
||||
monitor="val/acc_top1",
|
||||
mode="max"
|
||||
),
|
||||
|
||||
LearningRateMonitor(
|
||||
logging_interval="epoch"
|
||||
),
|
||||
|
||||
EarlyStopping(
|
||||
monitor="val/loss",
|
||||
min_delta=2e-4,
|
||||
patience=10,
|
||||
verbose=False,
|
||||
mode="min"
|
||||
)
|
||||
]
|
||||
|
||||
logger = MLFlowLogger(
|
||||
experiment_name="ImageNetMyAlgoritms",
|
||||
run_name="AlexNetCustom",
|
||||
tracking_uri="http://10.16.18.1:8055",
|
||||
log_model=True,
|
||||
)
|
||||
|
||||
alexnet = AlexNetCustom(in_channels=3, out_classes=1000)
|
||||
model_adam = ClassificationLightning(
|
||||
model=alexnet,
|
||||
num_classes=1000,
|
||||
optimizer_class=optimizer_Adam,
|
||||
optimizer_kwargs=optimizer_Adam_kwargs,
|
||||
scheduler_class=scheduler_ReduceLR,
|
||||
scheduler_kwargs=scheduler_ReduceLR_kwargs,
|
||||
scheduler_monitor="val/loss"
|
||||
)
|
||||
|
||||
trainer = pl.Trainer(
|
||||
max_epochs=num_epoch,
|
||||
accelerator="gpu",
|
||||
devices=1,
|
||||
logger=logger,
|
||||
# precision="16-mixed",
|
||||
callbacks=callbacks
|
||||
)
|
||||
|
||||
trainer.fit(model=model_adam,
|
||||
datamodule=imagenet_dm)
|
||||
@@ -1,244 +0,0 @@
|
||||
# %% Импорт библиотек
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from torchvision import datasets, transforms
|
||||
from torchmetrics.classification import Accuracy
|
||||
# %% LeeNet
|
||||
|
||||
class LeNetCustom(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_channels: int=1,
|
||||
out_classes: int=10):
|
||||
super().__init__()
|
||||
|
||||
self._out_classes=out_classes
|
||||
|
||||
self.feature = nn.Sequential(
|
||||
nn.Conv2d(in_channels, 6, kernel_size=5),
|
||||
# nn.Tanh(),
|
||||
nn.ReLU(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||||
|
||||
nn.Conv2d(6, 16, kernel_size=5),
|
||||
# nn.Tanh(),
|
||||
nn.ReLU(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2)
|
||||
)
|
||||
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(256, 120),
|
||||
# nn.LazyLinear(120),
|
||||
# nn.Tanh(),
|
||||
nn.ReLU(),
|
||||
nn.Linear(120, 84),
|
||||
# nn.Tanh(),
|
||||
nn.ReLU(),
|
||||
nn.Linear(84, out_classes)
|
||||
)
|
||||
|
||||
def forward(self, image):
|
||||
x = self.feature(image)
|
||||
x = self.mlp(x.flatten(1))
|
||||
|
||||
return x
|
||||
|
||||
class LeNetCustomLightning(pl.LightningModule):
|
||||
def __init__(self,
|
||||
model: nn.Module,
|
||||
lr_rate: float=0.1):
|
||||
super().__init__()
|
||||
|
||||
self.save_hyperparameters(ignore=["model"])
|
||||
self.model = model
|
||||
|
||||
self.criterion = nn.CrossEntropyLoss()
|
||||
self.train_metric = Accuracy(task="multiclass",
|
||||
num_classes=model._out_classes)
|
||||
self.val_metric = self.train_metric.clone()
|
||||
self.test_metric = self.train_metric.clone()
|
||||
|
||||
def forward(self, image):
|
||||
return self.model(image)
|
||||
|
||||
def _shared_step(self, sample, stage, metric):
|
||||
x, y = sample
|
||||
|
||||
logits = self.forward(x.to(torch.float32))
|
||||
preds = torch.argmax(logits, 1)
|
||||
loss = self.criterion(logits, y.to(torch.int64))
|
||||
|
||||
metric.update(preds, y)
|
||||
|
||||
self.log(f"{stage}_loss", loss, prog_bar=True,
|
||||
on_epoch=True, on_step=False)
|
||||
|
||||
return loss
|
||||
|
||||
def _log_metric(self, stage, metric):
|
||||
computed = metric.compute()
|
||||
|
||||
self.log(f"{stage}_acc", computed, prog_bar=True)
|
||||
|
||||
metric.reset()
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "train", self.train_metric)
|
||||
|
||||
def on_train_epoch_end(self):
|
||||
self._log_metric("train", self.train_metric)
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "val", self.val_metric)
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
self._log_metric("val", self.val_metric)
|
||||
|
||||
def test_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "test", self.test_metric)
|
||||
|
||||
def on_test_epoch_end(self):
|
||||
self._log_metric("test", self.test_metric)
|
||||
|
||||
def configure_optimizers(self):
|
||||
optimizer = torch.optim.SGD(self.parameters(),
|
||||
lr=self.hparams.lr_rate)
|
||||
|
||||
return optimizer
|
||||
|
||||
# %% Загрузка набора данных
|
||||
|
||||
batch_size = 64
|
||||
|
||||
image_transforms = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.1307,), (0.3081,))
|
||||
])
|
||||
|
||||
train_dataset = datasets.MNIST(
|
||||
root="./Datasets",
|
||||
train=True,
|
||||
download=False, # Если скачен, то убрать
|
||||
transform=image_transforms
|
||||
)
|
||||
|
||||
test_dataset = datasets.MNIST(
|
||||
root="./Datasets",
|
||||
train=False,
|
||||
download=False, # Если скачен, то убрать
|
||||
transform=image_transforms
|
||||
)
|
||||
|
||||
train_loader = DataLoader(dataset=train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True)
|
||||
|
||||
test_loader = DataLoader(dataset=test_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False)
|
||||
|
||||
# %% Запуск обучения
|
||||
|
||||
num_epoch=10
|
||||
|
||||
lenet = LeNetCustom(in_channels=1,
|
||||
out_classes=10)
|
||||
lenet_lightning = LeNetCustomLightning(model=lenet,
|
||||
lr_rate=0.1)
|
||||
lenet_lightning.compile()
|
||||
|
||||
trainer = pl.Trainer(max_epochs=num_epoch,
|
||||
accelerator="gpu",
|
||||
devices=1)
|
||||
|
||||
trainer.fit(model=lenet_lightning,
|
||||
train_dataloaders=train_loader,
|
||||
val_dataloaders=test_loader)
|
||||
|
||||
|
||||
|
||||
# %% Matlab Export
|
||||
|
||||
model = lenet_lightning.model
|
||||
model.eval()
|
||||
model = model.cpu()
|
||||
|
||||
dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
_ = model(dummy_input)
|
||||
|
||||
traced_model = torch.jit.trace(model, dummy_input)
|
||||
traced_model.save("../models/LeNetCustom.pt")
|
||||
|
||||
# %% ONNX экспорт
|
||||
|
||||
model = lenet_lightning.model.cpu().eval()
|
||||
|
||||
dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
|
||||
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy_input,
|
||||
"../models/LeNetCustom.onnx",
|
||||
input_names=["input"],
|
||||
output_names=["logits"],
|
||||
dynamic_axes={
|
||||
"input": {0: "batch_size"},
|
||||
"logits": {0: "batch_size"}
|
||||
},
|
||||
opset_version=17,
|
||||
# Проверка/оптимизация графа
|
||||
do_constant_folding=True
|
||||
)
|
||||
|
||||
# %% Новый экспорт в ONNX
|
||||
|
||||
model = lenet_lightning.model.cpu().eval()
|
||||
dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
|
||||
|
||||
onnx_program = torch.onnx.export(
|
||||
model,
|
||||
(dummy_input,),
|
||||
dynamo=True,
|
||||
verify=True
|
||||
)
|
||||
|
||||
onnx_program.save("../models/LeNetCustomNewExport.onnx")
|
||||
|
||||
# %% Инференс модели
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = torch.jit.load("./models/LeNetCustom.pt", map_location=device)
|
||||
model.eval()
|
||||
|
||||
image_transforms = transforms.Compose([
|
||||
transforms.Grayscale(),
|
||||
transforms.Resize((28,28)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.1307,), (0.3081,))
|
||||
])
|
||||
|
||||
def predict(image_path):
|
||||
image = Image.open(image_path)
|
||||
image = image_transforms(image)
|
||||
image = image.unsqueeze(0)
|
||||
image = image.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
output = model(image)
|
||||
|
||||
pred = output.argmax(dim=1).item()
|
||||
|
||||
return pred
|
||||
|
||||
print(f"Predicted label: {predict('./test_images/mnist/3.png')}")
|
||||
print(f"Predicted label: {predict('./test_images/mnist/5.png')}")
|
||||
print(f"Predicted label: {predict('./test_images/mnist/6.png')}")
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
from .lenet import LeNetCustom
|
||||
from .alexnet import AlexNetCustom
|
||||
from .vgg import VGGCustom
|
||||
|
||||
__all__ = [
|
||||
'LeNetCustom',
|
||||
'AlexNetCustom',
|
||||
'VGGCustom'
|
||||
]
|
||||
@@ -0,0 +1,53 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class AlexNetCustom(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels: int = 3,
|
||||
out_classes: int=1_000):
|
||||
super().__init__()
|
||||
|
||||
self._out_classes = out_classes
|
||||
|
||||
self.feature = nn.Sequential(
|
||||
nn.Conv2d(in_channels, 96, kernel_size=11, stride=4, padding=2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2),
|
||||
|
||||
nn.Conv2d(96, 256, kernel_size=5, padding=2),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2),
|
||||
|
||||
nn.Conv2d(256, 384, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(384, 384, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(384, 256, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(kernel_size=3, stride=2)
|
||||
)
|
||||
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Dropout(p=0.5),
|
||||
nn.Linear(256*6*6, 4096),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Dropout(p=0.5),
|
||||
nn.Linear(4096, 4096),
|
||||
nn.ReLU(inplace=True),
|
||||
|
||||
nn.Linear(4096, self._out_classes)
|
||||
)
|
||||
|
||||
def forward(self, image):
|
||||
x = self.feature(image)
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.mlp(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,38 @@
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class LeNetCustom(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_channels: int=1,
|
||||
out_classes: int=10):
|
||||
super().__init__()
|
||||
|
||||
self._out_classes=out_classes
|
||||
|
||||
self.feature = nn.Sequential(
|
||||
nn.Conv2d(in_channels, 6, kernel_size=5),
|
||||
nn.Tanh(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||||
|
||||
nn.Conv2d(6, 16, kernel_size=5),
|
||||
nn.Tanh(),
|
||||
nn.AvgPool2d(kernel_size=2, stride=2),
|
||||
|
||||
nn.Conv2d(16, 120, kernel_size=5),
|
||||
nn.Tanh(),
|
||||
)
|
||||
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.LazyLinear(84),
|
||||
nn.Tanh(),
|
||||
nn.Linear(84, out_classes)
|
||||
)
|
||||
|
||||
def forward(self, image):
|
||||
x = self.feature(image)
|
||||
x = self.mlp(x.flatten(1))
|
||||
|
||||
return x
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
VGG_CONFIG = {
|
||||
"VGG16": [64, 64, "M",
|
||||
128, 128, "M",
|
||||
256, 256, 256, "M",
|
||||
512, 512, 512, "M",
|
||||
512, 512, 512, "M"],
|
||||
"VGG19": [64, 64, "M",
|
||||
128, 128, "M",
|
||||
256, 256, 256, 256, "M",
|
||||
512, 512, 512, 512, "M",
|
||||
512, 512, 512, 512, "M"]
|
||||
}
|
||||
|
||||
|
||||
class VGGCustom(nn.Module):
|
||||
def __init__(self,
|
||||
architecture:str="VGG16",
|
||||
num_channels:int = 3,
|
||||
num_classes:int=10,
|
||||
batch_normalization=False,
|
||||
dropout:float=0.5) -> None:
|
||||
|
||||
super().__init__()
|
||||
|
||||
if architecture not in VGG_CONFIG:
|
||||
raise ValueError(f"Architecture {architecture} is invalid. ",
|
||||
f"Available architectures are :{list(VGG_CONFIG.keys())}")
|
||||
|
||||
|
||||
self.batch_normalization = batch_normalization
|
||||
|
||||
self.features = self._make_layers(
|
||||
VGG_CONFIG[architecture],
|
||||
in_channels=num_channels
|
||||
)
|
||||
|
||||
self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
|
||||
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(512*7*7, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(p=dropout),
|
||||
|
||||
nn.Linear(4096, 4096),
|
||||
nn.ReLU(True),
|
||||
nn.Dropout(p=dropout),
|
||||
|
||||
nn.Linear(4096, num_classes)
|
||||
)
|
||||
|
||||
def _make_layers(self, cfg, in_channels):
|
||||
layers = []
|
||||
|
||||
for layer in cfg:
|
||||
if layer == 'M':
|
||||
layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
|
||||
else:
|
||||
|
||||
layers.append(nn.Conv2d(in_channels=in_channels,
|
||||
out_channels=layer,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1))
|
||||
|
||||
if self.batch_normalization:
|
||||
layers.append(nn.BatchNorm2d(layer))
|
||||
|
||||
layers.append(nn.ReLU(True))
|
||||
|
||||
in_channels = layer
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, image):
|
||||
x = self.features(image)
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
x = self.classifier(x)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,7 @@
|
||||
from .imagenet import ImageNetDataModule, ImageNetValDataset
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ImageNetDataModule",
|
||||
"ImageNetValDataset",
|
||||
]
|
||||
@@ -0,0 +1,158 @@
|
||||
import torch
|
||||
import pandas as pd
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from PIL import Image
|
||||
from pathlib import Path
|
||||
from torchvision import datasets
|
||||
from torchvision.io import read_file, decode_image, ImageReadMode
|
||||
from typing import Optional, Callable, Any
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
|
||||
TransformType = Optional[Callable[..., Any]]
|
||||
|
||||
|
||||
def image_loader(path: str) -> torch.Tensor:
|
||||
data = read_file(path)
|
||||
|
||||
return decode_image(
|
||||
data,
|
||||
mode=ImageReadMode.RGB
|
||||
)
|
||||
|
||||
|
||||
class ImageNetValDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
image_dir,
|
||||
labels_file,
|
||||
class_to_idx,
|
||||
classes,
|
||||
transform=None
|
||||
):
|
||||
self.image_dir = Path(image_dir)
|
||||
self.transform = transform
|
||||
|
||||
self.class_to_idx = class_to_idx
|
||||
self.classes = classes
|
||||
|
||||
df = pd.read_csv(
|
||||
labels_file,
|
||||
usecols=["ImageId", "PredictionString"]
|
||||
)
|
||||
|
||||
image_ids = df["ImageId"].to_numpy()
|
||||
|
||||
synsets = (
|
||||
df["PredictionString"]
|
||||
.str.split()
|
||||
.str[0]
|
||||
.to_numpy()
|
||||
)
|
||||
|
||||
self.samples = [
|
||||
(
|
||||
self.image_dir / f"{image_id}.JPEG",
|
||||
class_to_idx[synset]
|
||||
)
|
||||
for image_id, synset in zip(image_ids, synsets)
|
||||
]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_path, label = self.samples[idx]
|
||||
|
||||
image = decode_image(
|
||||
str(image_path),
|
||||
mode=ImageReadMode.RGB
|
||||
)
|
||||
|
||||
if self.transform is not None:
|
||||
image = self.transform(image)
|
||||
|
||||
return image, label
|
||||
|
||||
|
||||
class ImageNetDataModule(pl.LightningDataModule):
|
||||
def __init__(self,
|
||||
dataset_root: Path,
|
||||
train_transform: TransformType,
|
||||
eval_transform: TransformType,
|
||||
batch_size: int = 64,
|
||||
num_workers:int = 8
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dataset_root = Path(dataset_root)
|
||||
|
||||
self.train_transform = train_transform
|
||||
self.eval_transform = eval_transform
|
||||
self.num_workers = num_workers
|
||||
self.batch_size = batch_size
|
||||
|
||||
self.imagenet_root = Path(dataset_root) / "ILSVRC"
|
||||
|
||||
self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
|
||||
self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
|
||||
|
||||
self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
|
||||
self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"
|
||||
|
||||
self.train_dataset = None
|
||||
self.val_dataset = None
|
||||
|
||||
self.class_to_idx = None
|
||||
self.classes = None
|
||||
|
||||
def setup(self, stage=None):
|
||||
synset_to_name = {}
|
||||
with open(self.mapping_file, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
synset, class_name = line.strip().split(" ", 1)
|
||||
synset_to_name[synset] = class_name
|
||||
|
||||
train_dataset = datasets.ImageFolder(root=self.train_dir,
|
||||
transform=self.train_transform,
|
||||
loader=image_loader)
|
||||
|
||||
self.class_to_idx = train_dataset.class_to_idx.copy()
|
||||
self.classes = [synset_to_name[synset]
|
||||
for synset in train_dataset.classes]
|
||||
train_dataset.classes = self.classes
|
||||
self.train_dataset = train_dataset
|
||||
|
||||
self.val_dataset = ImageNetValDataset(
|
||||
image_dir=self.val_dir,
|
||||
labels_file=self.val_labels_file,
|
||||
class_to_idx=self.class_to_idx,
|
||||
classes=self.classes,
|
||||
transform=self.eval_transform
|
||||
)
|
||||
|
||||
def train_dataloader(self):
|
||||
return DataLoader(
|
||||
dataset=self.train_dataset,
|
||||
batch_size=self.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=self.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=self.num_workers > 0,
|
||||
prefetch_factor=2 if self.num_workers > 0 else None,
|
||||
drop_last=True
|
||||
)
|
||||
|
||||
def val_dataloader(self):
|
||||
return DataLoader(
|
||||
dataset=self.val_dataset,
|
||||
batch_size=self.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=self.num_workers,
|
||||
pin_memory=True,
|
||||
persistent_workers=self.num_workers > 0,
|
||||
prefetch_factor=2 if self.num_workers > 0 else None,
|
||||
drop_last=False
|
||||
)
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import os
|
||||
import torch
|
||||
import mlflow
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from torch import optim
|
||||
from dotenv import load_dotenv
|
||||
from torchvision.transforms import v2
|
||||
from datasets import ImageNetDataModule
|
||||
from train import ClassificationLightning
|
||||
from custom_arhitecture import VGGCustom
|
||||
|
||||
from pytorch_lightning.loggers import MLFlowLogger
|
||||
from pytorch_lightning.callbacks import LearningRateMonitor
|
||||
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
|
||||
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
|
||||
|
||||
|
||||
load_dotenv()
|
||||
|
||||
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
||||
IMAGENET_STD = [0.229, 0.224, 0.225]
|
||||
|
||||
dataset_root = "/media/daniel/developments/Datasets/ImageNet"
|
||||
batch_size = 48
|
||||
num_workers = 8
|
||||
num_epoch = 90
|
||||
learning_rate = 3e-4
|
||||
|
||||
train_transform = v2.Compose([
|
||||
v2.RandomResizedCrop(
|
||||
size=(224, 224),
|
||||
scale=(0.08, 1.0),
|
||||
ratio=(3 / 4, 4 / 3),
|
||||
antialias=True
|
||||
),
|
||||
v2.RandomHorizontalFlip(p=0.5),
|
||||
v2.ToDtype(torch.float32, scale=True),
|
||||
v2.Normalize(
|
||||
mean=IMAGENET_MEAN,
|
||||
std=IMAGENET_STD
|
||||
)
|
||||
])
|
||||
|
||||
eval_transform = v2.Compose([
|
||||
v2.Resize(256, antialias=True),
|
||||
v2.CenterCrop(224),
|
||||
v2.ToDtype(torch.float32, scale=True),
|
||||
v2.Normalize(
|
||||
mean=IMAGENET_MEAN,
|
||||
std=IMAGENET_STD
|
||||
)
|
||||
])
|
||||
|
||||
|
||||
pl.seed_everything(42, workers=True)
|
||||
torch.set_float32_matmul_precision("high")
|
||||
|
||||
imagenet_dm = ImageNetDataModule(
|
||||
dataset_root=dataset_root,
|
||||
train_transform=train_transform,
|
||||
eval_transform=eval_transform,
|
||||
batch_size=batch_size,
|
||||
num_workers=num_workers
|
||||
)
|
||||
|
||||
optimizer_Adam_kwargs = {
|
||||
"lr": 3e-4,
|
||||
"weight_decay": 1e-4
|
||||
}
|
||||
|
||||
scheduler_ReduceLR_kwargs={
|
||||
"mode": "min",
|
||||
"factor": 0.1,
|
||||
"patience": 5
|
||||
}
|
||||
|
||||
callbacks = [
|
||||
ModelCheckpoint(
|
||||
dirpath="models",
|
||||
filename="epoch={epoch}",
|
||||
auto_insert_metric_name=False,
|
||||
save_last=True,
|
||||
save_top_k=1,
|
||||
monitor="val/acc_top1",
|
||||
mode="max"
|
||||
),
|
||||
|
||||
LearningRateMonitor(
|
||||
logging_interval="epoch"
|
||||
),
|
||||
|
||||
EarlyStopping(
|
||||
monitor="val/loss",
|
||||
min_delta=2e-4,
|
||||
patience=10,
|
||||
verbose=False,
|
||||
mode="min"
|
||||
)
|
||||
]
|
||||
|
||||
# Системные метрики
|
||||
mlflow.enable_system_metrics_logging()
|
||||
mlflow.set_system_metrics_sampling_interval(5)
|
||||
|
||||
logger = MLFlowLogger(
|
||||
experiment_name=os.getenv("EXPERIMENT_NAME"),
|
||||
run_name="VGG16Custom",
|
||||
tracking_uri=os.getenv("TRACKING_URI"),
|
||||
log_model=True,
|
||||
)
|
||||
|
||||
|
||||
vgg = VGGCustom(architecture="VGG16",
|
||||
num_channels=3,
|
||||
num_classes=1000,
|
||||
batch_normalization=False,
|
||||
dropout=.5)
|
||||
|
||||
# Необходимо чтобы отрпаботал LazyLinear и принял размерности!
|
||||
# with torch.no_grad():
|
||||
# alexnet(torch.zeros(1, 3, 224, 224))
|
||||
|
||||
model = ClassificationLightning(model=vgg,
|
||||
num_classes=1000,
|
||||
optimizer_class=optim.Adam,
|
||||
optimizer_kwargs=optimizer_Adam_kwargs,
|
||||
scheduler_class=optim.lr_scheduler.ReduceLROnPlateau,
|
||||
scheduler_kwargs=scheduler_ReduceLR_kwargs,
|
||||
scheduler_monitor="val/loss")
|
||||
|
||||
trainer = pl.Trainer(
|
||||
max_epochs=num_epoch,
|
||||
accelerator='gpu',
|
||||
devices=1,
|
||||
precision="16-mixed",
|
||||
logger=logger,
|
||||
callbacks= callbacks
|
||||
)
|
||||
|
||||
trainer.fit(model=model,
|
||||
datamodule=imagenet_dm)
|
||||
@@ -0,0 +1,6 @@
|
||||
from .classification import ClassificationLightning
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ClassificationLightning"
|
||||
]
|
||||
@@ -0,0 +1,165 @@
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from torch import optim
|
||||
from torchmetrics import MetricCollection
|
||||
from torchmetrics.classification import Accuracy
|
||||
|
||||
|
||||
class ClassificationLightning(pl.LightningModule):
|
||||
def __init__(self,
|
||||
model: nn.Module,
|
||||
num_classes: int,
|
||||
optimizer_class,
|
||||
optimizer_kwargs: dict,
|
||||
scheduler_class=None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
scheduler_monitor: str = "val/loss"):
|
||||
super().__init__()
|
||||
|
||||
|
||||
scheduler_kwargs = scheduler_kwargs or {}
|
||||
|
||||
self.model = model
|
||||
self.optimizer_class = optimizer_class
|
||||
self.scheduler_class = scheduler_class
|
||||
|
||||
num_parameters = sum(
|
||||
p.numel()
|
||||
for p in model.parameters()
|
||||
)
|
||||
|
||||
num_trainable_parameters = sum(
|
||||
p.numel()
|
||||
for p in model.parameters()
|
||||
if p.requires_grad
|
||||
)
|
||||
|
||||
self.save_hyperparameters({
|
||||
"model": model.__class__.__name__,
|
||||
"num_classes": num_classes,
|
||||
|
||||
"num_parameters": num_parameters,
|
||||
"num_trainable_parameters": num_trainable_parameters,
|
||||
|
||||
"optimizer": optimizer_class.__name__,
|
||||
"optimizer_kwargs": optimizer_kwargs,
|
||||
|
||||
"scheduler": (
|
||||
scheduler_class.__name__
|
||||
if scheduler_class is not None
|
||||
else None
|
||||
),
|
||||
"scheduler_kwargs": scheduler_kwargs,
|
||||
"scheduler_monitor": scheduler_monitor
|
||||
})
|
||||
|
||||
|
||||
self.criterion = nn.CrossEntropyLoss()
|
||||
self.train_metrics = MetricCollection({
|
||||
"acc_top1": Accuracy(task="multiclass",
|
||||
num_classes=self.hparams.num_classes,
|
||||
top_k=1),
|
||||
"acc_top5": Accuracy(task="multiclass",
|
||||
num_classes=self.hparams.num_classes,
|
||||
top_k=5),
|
||||
})
|
||||
|
||||
self.val_metrics = self.train_metrics.clone()
|
||||
self.test_metrics = self.train_metrics.clone()
|
||||
|
||||
def forward(self, image):
|
||||
return self.model(image)
|
||||
|
||||
def _shared_step(self,
|
||||
sample,
|
||||
stage,
|
||||
metrics):
|
||||
x, y = sample
|
||||
|
||||
logits = self(x)
|
||||
loss = self.criterion(logits, y)
|
||||
|
||||
metrics.update(logits, y)
|
||||
|
||||
self.log(f"{stage}/loss",
|
||||
loss,
|
||||
prog_bar=True,
|
||||
on_epoch=True,
|
||||
on_step=False
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
def _log_metrics(self, stage, metrics):
|
||||
computed = metrics.compute()
|
||||
|
||||
for name, value in computed.items():
|
||||
self.log(f"{stage}/{name}", value, prog_bar=True)
|
||||
|
||||
metrics.reset()
|
||||
|
||||
def on_train_epoch_start(self):
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self._epoch_start_time = time.perf_counter()
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "train", self.train_metrics)
|
||||
|
||||
def on_train_epoch_end(self):
|
||||
self._log_metrics("train", self.train_metrics)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
epoch_time = time.perf_counter() - self._epoch_start_time
|
||||
lr = self.optimizers().param_groups[0]["lr"]
|
||||
|
||||
self.log("train/lr", lr, prog_bar=True)
|
||||
self.log("train/epoch_time", epoch_time)
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "val", self.val_metrics)
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
self._log_metrics("val", self.val_metrics)
|
||||
|
||||
def test_step(self, batch, batch_idx):
|
||||
return self._shared_step(batch, "test", self.test_metrics)
|
||||
|
||||
def on_test_epoch_end(self):
|
||||
self._log_metrics("test", self.test_metrics)
|
||||
|
||||
def configure_optimizers(self):
|
||||
optimizer = self.optimizer_class(
|
||||
self.parameters(),
|
||||
**self.hparams.optimizer_kwargs
|
||||
)
|
||||
|
||||
if self.scheduler_class is None:
|
||||
return optimizer
|
||||
|
||||
scheduler = self.scheduler_class(
|
||||
optimizer,
|
||||
**self.hparams.scheduler_kwargs
|
||||
)
|
||||
|
||||
scheduler_config = {
|
||||
"scheduler": scheduler,
|
||||
"interval": "epoch",
|
||||
"frequency": 1
|
||||
}
|
||||
|
||||
if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
|
||||
scheduler_config["monitor"] = self.hparams.scheduler_monitor
|
||||
|
||||
return {
|
||||
"optimizer": optimizer,
|
||||
"lr_scheduler": scheduler_config
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user