train.py

import os
import torch
from torch.utils.data import DataLoader
from datasets.dataset import PositioningDataset
from models.model import PositioningModel
from utils.transforms import preprocess_image
from utils.visualization import show_sample

# --- Настройки ---
BATCH_SIZE = 4
LEARNING_RATE = 0.001
NUM_EPOCHS = 20

# --- Пути к данным ---
current_dir = os.path.dirname(os.path.abspath(__file__))
excel_path = os.path.join(current_dir, "data", "choords.xlsx")
images_root = os.path.join(current_dir, "data", "images")

# --- Создание датасета ---
dataset = PositioningDataset(images_root)

# --- Создание DataLoader ---
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

# --- Модель ---
model = PositioningModel()

# --- Оптимизатор и функция потерь ---
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = torch.nn.MSELoss()

# --- Обучение модели ---
for epoch in range(NUM_EPOCHS):
    print(f"Epoch {epoch + 1}/{NUM_EPOCHS}")
    
    for batch_idx, (images, labels) in enumerate(dataloader):
        # Перемещаем данные на GPU, если доступен
        if torch.cuda.is_available():
            images = images.cuda()
            labels = labels.cuda()
        
        # Обнуляем градиенты
        optimizer.zero_grad()
        
        # Прямое распространение
        outputs = model(images)
        
        # Вычисляем ошибку
        loss = criterion(outputs, labels)
        
        # Backward pass
        loss.backward()
        
        # Обновляем параметры
        optimizer.step()
        
        # Выводим прогресс
        if batch_idx % 10 == 0:
            print(f"Batch {batch_idx}, Loss: {loss.item():.4f}")

    # Визуализация примера после каждой эпохи
    if torch.cuda.is_available():
        images = images.cpu()
        labels = labels.cpu()
        outputs = outputs.cpu()
    
    show_sample(images[0], labels[0], outputs[0])

print("Обучение завершено!")