grabardm/PyTorchCource
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Обнавлена структура проекта

Browse Source
This commit is contained in:
Даниил Грабарь
2026-09-28 15:40:05 +10:00
parent 53b9c6d854
commit 056dd87746
17 changed files with 671 additions and 733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+2
View File
@@ -0,0 +1,2 @@
EXPERIMENT_NAME=
TRACKING_URI=
.gitignore
+7 -1
View File
@@ -1,4 +1,10 @@
.vscode/
__pycache__/
models/ models/
Datasets/ Datasets/
test_images/ test_images/
train_logs/ train_logs/
lightning_logs/
.env
BaseCNNArhitecture/AlexNetCustom.py
-488
View File
@@ -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)
BaseCNNArhitecture/LeNetCustom.py
-244
View File
@@ -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')}")
custom_arhitecture/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .lenet import LeNetCustom
from .alexnet import AlexNetCustom
from .vgg import VGGCustom
__all__ = [
'LeNetCustom',
'AlexNetCustom',
'VGGCustom'
]
custom_arhitecture/alexnet.py
+53
View File
@@ -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
custom_arhitecture/lenet.py
+38
View File
@@ -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
custom_arhitecture/vgg.py
+84
View File
@@ -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
datasets/__init__.py
+7
View File
@@ -0,0 +1,7 @@
from .imagenet import ImageNetDataModule, ImageNetValDataset
__all__ = [
"ImageNetDataModule",
"ImageNetValDataset",
]
datasets/imagenet.py
+158
View File
@@ -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
)
MyDataExamplePL.py → exampels/MyDataExamplePL.py
View File
MyNNExamplePL.py → exampels/MyNNExamplePL.py
View File
SimpleMLP.py → exampels/SimpleMLP.py
View File
ffn.py → exampels/ffn.py
View File
main.py
+142
View File
@@ -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)
train/__init__.py
+6
View File
@@ -0,0 +1,6 @@
from .classification import ClassificationLightning
__all__ = [
"ClassificationLightning"
]
train/classification.py
+165
View File
@@ -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
}