Added project

train.py

@@ -0,0 +1,48 @@
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from datasets import get_data_loaders
+from model import SimpleCNN
+from utils import save_model, calculate_accuracy
+
+# Настройки
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = {
+    "batch_size": 128,
+    "epochs": 10,
+    "learning_rate": 0.001,
+    "data_dir": "./data",
+    "logs_dir": "./logs",
+}
+
+# Загрузка данных
+train_loader, test_loader = get_data_loaders(config["data_dir"], config["batch_size"])
+
+# Модель, критерий, оптимизатор
+model = SimpleCNN().to(device)
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=config["learning_rate"])
+
+# Тренировка модели
+for epoch in range(config["epochs"]):
+    model.train()
+    running_loss = 0.0
+    for inputs, labels in train_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+
+    print(f"Epoch {epoch+1}/{config['epochs']}, Loss: {running_loss/len(train_loader)}")
+
+# Тестирование
+accuracy = calculate_accuracy(model, test_loader, device)
+print(f"Test Accuracy: {accuracy:.2f}%")
+
+# Сохранение модели
+save_model(model, "./model.pth")