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

Добавлена реализация LeNet и Alexnet

Browse Source
This commit is contained in:
Даниил Грабарь
2026-09-17 15:05:30 +10:00
parent 41c9010db2
commit d689368d75
3 changed files with 603 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BaseCNNArhitecture/AlexNetCustom.py
+329
View File
@@ -0,0 +1,329 @@
# %% Импорт библиотек
import time
import torch
import kagglehub
import torch.nn as nn
import pytorch_lightning as pl
from torch import optim
from torchmetrics import MetricCollection
from torchvision import datasets, transforms
from torchmetrics.classification import Accuracy
from torch.utils.data import DataLoader, random_split
# автоматическое отслеживание 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 = "./Datasets/ImageNet"
batch_size = 64
num_workers = 8
num_epoch = 2
# %% Трансформации
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)
# %% Датасет
full_train_dataset = datasets.ImageNet(
root=dataset_root,
split="train",
transform=train_transform
)
train_size = int(0.9 * len(full_train_dataset))
val_size = len(full_train_dataset) - train_size
train_dataset, val_dataset = random_split(
full_train_dataset,
[train_size, val_size],
generator=torch.Generator().manual_seed(42)
)
val_base_dataset = datasets.ImageNet(
root=dataset_root,
split="train",
transform=eval_transform
)
val_dataset.dataset = val_base_dataset
test_dataset = datasets.ImageNet(
root=dataset_root,
split="val",
transform=eval_transform
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
persistent_workers=True
)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
persistent_workers=True
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
persistent_workers=True
)
# %% 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
x = x.to(torch.float32)
y = y.to(torch.int64)
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
}
BaseCNNArhitecture/LeNetCustom.py
+244
View File
@@ -0,0 +1,244 @@
# %% Импорт библиотек
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')}")
MyNNExamplePL.py
+28 -1
View File
@@ -178,6 +178,8 @@ test_loader = DataLoader(dataset=test_dataset,
# %% Запуск обучения
model = MyNNExample(model=MyTorchCNNModel())
model.compile()
callbacks = [
ModelCheckpoint(dirpath="models",
filename="{epoch}_{val_accuracy:.4f}_{val_loss:.4f}",
@@ -209,7 +211,7 @@ trainer.fit(model=model,
# %% Инференс модели
checkpoint_path = "./models/6_0.9460_0.1643.ckpt"
checkpoint_path = "./models/7_0.9720_0.0958.ckpt"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyNNExample.load_from_checkpoint(checkpoint_path,
model=MyTorchCNNModel(output_size=10))
@@ -227,3 +229,28 @@ with torch.no_grad():
print(f"Предсказание: {pred_label}")
print(f"Истинная метка: {true_label}")
print(f"Уверенность: {probs[0, pred_label].item():.2%}")
# %% Matlab export
pytorch_model = model.model
pytorch_model.eval()
pytorch_model = pytorch_model.cpu()
dummy_input = torch.randn(1, 1, 28, 28, dtype=torch.float32)
with torch.no_grad():
_ = pytorch_model(dummy_input)
traced_model = torch.jit.trace(pytorch_model, dummy_input)
with torch.no_grad():
y_original = pytorch_model(dummy_input)
y_traced = traced_model(dummy_input)
print("Original:", y_original.shape)
print("Traced:", y_traced.shape)
print("Max difference:", torch.max(torch.abs(y_original - y_traced)).item())
traced_model.save("./models/my_model_matlab.pt")