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

2026-09-28 15:40:05 +10:00
parent 53b9c6d854
commit 056dd87746
17 changed files with 671 additions and 733 deletions
@@ -0,0 +1,189 @@
+# %% Импорт библиотек
+
+import tqdm
+import torch
+import torch.nn as nn
+
+from torch.utils.data import DataLoader
+from torchmetrics import MetricCollection
+from torchvision import datasets, transforms
+from torchmetrics.classification import Accuracy, Precision, Recall, F1Score
+
+from PIL import Image
+from beartype import beartype
+
+# %% Гиперпараметры
+num_epoch = 8
+batch_size = 64
+learning_rate = 1e-3
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# %% Трансформация данных
+
+# Resize(size) — изменяет размер изображения.
+# Grayscale() — преобразование изображения в оттенки серого.
+# ToTensor() — преобразует изображение PIL или NumPy-массив в PyTorch-тензор.
+# Normalize(mean, std) — нормализация тензора изображения.
+# GaussianBlur(kernel_size) — размытие по Гауссу.
+
+
+image_transforms = transforms.Compose([
+    transforms.ToTensor()
+])
+
+# %% Загрузка набора данных
+
+train_dataset = datasets.MNIST(
+    root="./MNIST",
+    train=True,
+    download=False, # Если скачен, то убрать
+    transform=image_transforms
+)
+
+test_dataset = datasets.MNIST(
+    root="./MNIST",
+    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)
+# %% Структура модели
+
+class FFN(nn.Module):
+    
+    @beartype
+    def __init__(self,
+                 input_size: int,
+                 output_size: int):
+        super().__init__()
+        
+        self.input_layer = nn.Linear(input_size, 32*32)
+        self.hidden_layer_1 = nn.Linear(32*32, 32*32)
+        self.hidden_layer_2 = nn.Linear(32*32, 32*32)
+        self.output_layer = nn.Linear(32*32, output_size)
+    
+    @beartype
+    def forward(self, image: torch.Tensor) -> torch.Tensor:
+        x = image.flatten(start_dim=1)  # batch_size x 28*28
+        x = torch.relu(self.input_layer(x))
+        x = torch.relu(self.hidden_layer_1(x))
+        x = torch.relu(self.hidden_layer_2(x))
+        x = self.output_layer(x)
+        
+        return x # batch_size x 10
+
+model = FFN(input_size=28*28,
+            output_size=10)
+model.to(device)
+
+# %% Метрики
+
+train_metrics = MetricCollection({
+    "accuracy": Accuracy(task="multiclass", num_classes=10),
+    "precision": Precision(task="multiclass", num_classes=10, average="macro"),    
+    "recall": Recall(task="multiclass", num_classes=10, average="macro"),
+    "f1": F1Score(task="multiclass", num_classes=10, average="macro")    
+}).to(device)
+
+test_metrics = MetricCollection({
+    "accuracy": Accuracy(task="multiclass", num_classes=10),
+    "precision": Precision(task="multiclass", num_classes=10, average="macro"),    
+    "recall": Recall(task="multiclass", num_classes=10, average="macro"),
+    "f1": F1Score(task="multiclass", num_classes=10, average="macro")    
+}).to(device)
+
+# %% Оптимизаторы
+
+criterion = nn.CrossEntropyLoss()
+optim = torch.optim.Adam(model.parameters(), lr=learning_rate)
+# %% Тренировочный цикл
+
+for epoch in tqdm.tqdm(range(num_epoch)):
+    model.train()
+    running_loss = 0
+    
+    for data, target in train_loader:
+        data, target = data.to(device), target.to(device)
+        
+        optim.zero_grad()
+        outputs = model(data)
+        loss = criterion(outputs, target)
+        loss.backward()
+        running_loss += loss.item()
+        optim.step()
+        
+        preds = outputs.argmax(dim=1)
+        train_metrics.update(preds, target)
+    
+    results = train_metrics.compute()        
+    print(f"Epoch: {epoch+1}")
+    print("=== TRAIN ===")
+    print(f"  Loss: {(running_loss/len(train_loader)):.4f}")
+    print(f"  Accuracy:  {results['accuracy']:.4f}")
+    print(f"  Precision: {results['precision']:.4f}")
+    print(f"  Recall:    {results['recall']:.4f}")
+    print(f"  F1:        {results['f1']:.4f}")
+    print()
+    
+    train_metrics.reset()
+    
+    model.eval()
+    with torch.no_grad():
+        test_loss = 0
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            
+            outputs = model(data)
+            loss = criterion(outputs, target)
+            test_loss += loss.item()
+            preds = outputs.argmax(dim=1)
+            test_metrics.update(preds, target)
+        
+        results = test_metrics.compute()        
+        
+        print("=== TEST ===")
+        print(f"  Loss: {(test_loss/len(train_loader)):.4f}")
+        print(f"  Accuracy:  {results['accuracy']:.4f}")
+        print(f"  Precision: {results['precision']:.4f}")
+        print(f"  Recall:    {results['recall']:.4f}")
+        print(f"  F1:        {results['f1']:.4f}")
+        print()
+        
+    test_metrics.reset()
+        
+    
+torch.save(model.state_dict(), "./models/mnist_model.pth")
+# %% Инференс модели
+
+image_transforms = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize((28,28)),
+    transforms.ToTensor()    
+])
+
+model.eval()
+model.load_state_dict(torch.load("./models/mnist_model.pth"))
+
+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')}")