Added project

This commit is contained in:
2025-01-13 11:18:03 +03:00
parent ec8e1c49de
commit 45150662a3
9 changed files with 175 additions and 0 deletions
+9
View File
@@ -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"]
+35
View File
@@ -0,0 +1,35 @@
# Инструкция по запуску
- Запускаем контейнер и подключаемся к нему
- Устанавливаем проект (тащим с гита, переносим любым удобным способом)
- Включаем venv
- Дотягиваем зависимости
- Запускаем как `python3 train.py`
## train.py — Основной скрипт для обучения модели
Этот файл запускает процесс обучения, тестирования и, возможно, сохранения модели. Он объединяет остальные компоненты проекта.
Содержимое:
- Импорты всех необходимых модулей (модель, данные, утилиты, конфигурации).
- Чтение гиперпараметров из config.yaml.
- Основной цикл обучения.
- Тестирование модели на валидационных данных.
## model.py — Определение архитектуры модели
Этот файл содержит только архитектуру модели. Если необходимо сменить модель, достаточно будет изменить только этот файл.
Содержимое:
- Определение слоёв и функции forward для прохождения данных через модель.
## datasets.py — Работа с данными
Этот файл отвечает за загрузку, подготовку и трансформацию данных.
Содержимое:
- Загрузка тренировочных и тестовых наборов данных.
- Определение необходимых трансформаций.
## utils.py — Вспомогательные функции
Этот файл содержит функции, которые помогают сократить основной код и избежать дублирования.
Содержимое:
- Сохранение/загрузка модели.
- Подсчёт точности.
+5
View File
@@ -0,0 +1,5 @@
batch_size: 128
epochs: 10
learning_rate: 0.001
data_dir: "./data"
logs_dir: "./logs"
+19
View File
@@ -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
+11
View File
@@ -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
+19
View File
@@ -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
+5
View File
@@ -0,0 +1,5 @@
torch==2.0.1
torchvision==0.15.2
matplotlib==3.7.2
pyyaml==6.0
numpy==1.24.1
+48
View File
@@ -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")
+24
View File
@@ -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