# %% Импорт библиотек import torch import torch.nn as nn import pytorch_lightning as pl from torch.utils.data import DataLoader from torchmetrics import MetricCollection from torchvision import datasets, transforms from torchmetrics.classification import Accuracy, Precision, Recall, F1Score # Автоматическое логирование результатов from pytorch_lightning.loggers import WandbLogger, TensorBoardLogger, CSVLogger # сохранение весов from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint # автоматическое отслеживание lr from pytorch_lightning.callbacks import LearningRateMonitor # ранние остановки from pytorch_lightning.callbacks.early_stopping import EarlyStopping # %% Гиперпараметры num_epoch = 8 batch_size = 64 # %% Создание модели PyTorch class MyTorchCNNModel(nn.Module): def __init__(self, output_size: int = 10): super().__init__() self.output_size = output_size self.feature = nn.Sequential( nn.LazyConv2d(64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(256), nn.ReLU(), nn.MaxPool2d(kernel_size=2) ) self.head = nn.Conv2d(256, self.output_size, kernel_size=1) self.pool = nn.AdaptiveAvgPool2d((1, 1)) def forward(self, image): x = self.feature(image) x = self.head(x) x = self.pool(x) return x.flatten(1) # %% Создание модели PyTorch Lightning class MyNNExample(pl.LightningModule): def __init__(self, model: nn.Module, eta=3e-4): super().__init__() self.save_hyperparameters(ignore=["model"]) # device = "cuda" if torch.cuda.is_available() else "cpu" self.model = model self.criterion = nn.CrossEntropyLoss() self.train_metrics = MetricCollection({ "accuracy": Accuracy(task="multiclass", num_classes=model.output_size), "precision": Precision(task="multiclass", num_classes=model.output_size, average="macro"), "recall": Recall(task="multiclass", num_classes=model.output_size, average="macro"), "f1": F1Score(task="multiclass", num_classes=model.output_size, average="macro") }) 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.forward(x.to(torch.float32)) preds = torch.argmax(logits, 1) loss = self.criterion(logits, y.to(torch.int64)) metrics.update(preds, 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 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) 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 = torch.optim.Adam(self.parameters(), lr=self.hparams.eta) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, patience=5, mode="min" ) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "monitor": "val_loss", "interval": "epoch", } } # %% Загрузка набора данных image_transforms = transforms.Compose([ transforms.ToTensor() ]) 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) # %% Запуск обучения model = MyNNExample(model=MyTorchCNNModel()) model.compile() callbacks = [ ModelCheckpoint(dirpath="models", filename="{epoch}_{val_accuracy:.4f}_{val_loss:.4f}", auto_insert_metric_name=False, save_last=True, save_top_k=1, monitor="val_loss", mode="min"), LearningRateMonitor(logging_interval="step"), EarlyStopping(monitor="val_loss", min_delta=2e-4, patience=10, verbose=False, mode="min") ] trainer = pl.Trainer( max_epochs=num_epoch, accelerator="gpu", devices=1, logger=CSVLogger("./train_logs", name="mnist_csv", version=None, flush_logs_every_n_steps=100), callbacks=callbacks ) trainer.fit(model=model, train_dataloaders=train_loader, val_dataloaders=test_loader) # %% Инференс модели 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)) model = model.to(device) model.eval() image, true_label = test_dataset[42] image_batch = image.unsqueeze(0).to(device) with torch.no_grad(): logits = model(image_batch) probs = torch.softmax(logits, dim=1) pred_label = logits.argmax(dim=1).item() 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")