diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ea658b7 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,9 @@ +FROM nvcr.io/nvidia/driver:550.127.08-ubuntu22.04 + +RUN apt-get update && apt-get install -y python3 python3-pip +WORKDIR /app + +COPY . . +RUN pip install -r requirements.txt + +CMD ["python3", "train.py"] diff --git a/README.md b/README.md index e69de29..c61cc53 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,35 @@ +# Инструкция по запуску +- Запускаем контейнер и подключаемся к нему +- Устанавливаем проект (тащим с гита, переносим любым удобным способом) +- Включаем venv +- Дотягиваем зависимости +- Запускаем как `python3 train.py` + +## train.py — Основной скрипт для обучения модели +Этот файл запускает процесс обучения, тестирования и, возможно, сохранения модели. Он объединяет остальные компоненты проекта. + +Содержимое: +- Импорты всех необходимых модулей (модель, данные, утилиты, конфигурации). +- Чтение гиперпараметров из config.yaml. +- Основной цикл обучения. +- Тестирование модели на валидационных данных. + +## model.py — Определение архитектуры модели +Этот файл содержит только архитектуру модели. Если необходимо сменить модель, достаточно будет изменить только этот файл. + +Содержимое: +- Определение слоёв и функции forward для прохождения данных через модель. + +## datasets.py — Работа с данными +Этот файл отвечает за загрузку, подготовку и трансформацию данных. + +Содержимое: +- Загрузка тренировочных и тестовых наборов данных. +- Определение необходимых трансформаций. + +## utils.py — Вспомогательные функции +Этот файл содержит функции, которые помогают сократить основной код и избежать дублирования. + +Содержимое: +- Сохранение/загрузка модели. +- Подсчёт точности. diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..eaf78fd --- /dev/null +++ b/config.yaml @@ -0,0 +1,5 @@ +batch_size: 128 +epochs: 10 +learning_rate: 0.001 +data_dir: "./data" +logs_dir: "./logs" diff --git a/datasets.py b/datasets.py new file mode 100644 index 0000000..ce70dc5 --- /dev/null +++ b/datasets.py @@ -0,0 +1,19 @@ +import torchvision +import torchvision.transforms as transforms +import torch.utils.data + +def get_data_loaders(data_dir, batch_size): + transform = transforms.Compose([ + transforms.RandomHorizontalFlip(), + transforms.RandomCrop(32, padding=4), + transforms.ToTensor(), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + ]) + + train_dataset = torchvision.datasets.CIFAR10(root=data_dir, train=True, download=True, transform=transform) + train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4) + + test_dataset = torchvision.datasets.CIFAR10(root=data_dir, train=False, download=True, transform=transform) + test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4) + + return train_loader, test_loader diff --git a/k8s/pod.yaml b/k8s/pod.yaml new file mode 100644 index 0000000..deee999 --- /dev/null +++ b/k8s/pod.yaml @@ -0,0 +1,11 @@ +apiVersion: v1 +kind: Pod +metadata: + name: pytorch-cifar10 +spec: + containers: + - name: pytorch-cifar10 + image: drn.tst.nubes.ru/gpu-test-pytorch-cifar:0.0.1 + resources: + limits: + nvidia.com/gpu: 1 diff --git a/model.py b/model.py new file mode 100644 index 0000000..eddb5f8 --- /dev/null +++ b/model.py @@ -0,0 +1,19 @@ +import torch.nn as nn +import torch.nn.functional as F + +class SimpleCNN(nn.Module): + def __init__(self): + super(SimpleCNN, self).__init__() + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool = nn.MaxPool2d(2, 2) + self.fc1 = nn.Linear(64 * 8 * 8, 128) + self.fc2 = nn.Linear(128, 10) + + def forward(self, x): + x = self.pool(F.relu(self.conv1(x))) + x = self.pool(F.relu(self.conv2(x))) + x = x.view(-1, 64 * 8 * 8) + x = F.relu(self.fc1(x)) + x = self.fc2(x) + return x diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..6db42a3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +torch==2.0.1 +torchvision==0.15.2 +matplotlib==3.7.2 +pyyaml==6.0 +numpy==1.24.1 diff --git a/train.py b/train.py new file mode 100644 index 0000000..00002b8 --- /dev/null +++ b/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") diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..888bf3d --- /dev/null +++ b/utils.py @@ -0,0 +1,24 @@ +import torch + +def save_model(model, path): + """Сохраняет модель в указанный файл.""" + torch.save(model.state_dict(), path) + +def load_model(model, path): + """Загружает веса модели из файла.""" + model.load_state_dict(torch.load(path)) + return model + +def calculate_accuracy(model, data_loader, device): + """Подсчитывает точность модели на заданном наборе данных.""" + model.eval() + correct = 0 + total = 0 + with torch.no_grad(): + for inputs, labels in data_loader: + inputs, labels = inputs.to(device), labels.to(device) + outputs = model(inputs) + _, predicted = torch.max(outputs, 1) + total += labels.size(0) + correct += (predicted == labels).sum().item() + return 100 * correct / total