Added project
This commit is contained in:
@@ -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"]
|
||||
@@ -0,0 +1,35 @@
|
||||
# Инструкция по запуску
|
||||
- Запускаем контейнер и подключаемся к нему
|
||||
- Устанавливаем проект (тащим с гита, переносим любым удобным способом)
|
||||
- Включаем venv
|
||||
- Дотягиваем зависимости
|
||||
- Запускаем как `python3 train.py`
|
||||
|
||||
## train.py — Основной скрипт для обучения модели
|
||||
Этот файл запускает процесс обучения, тестирования и, возможно, сохранения модели. Он объединяет остальные компоненты проекта.
|
||||
|
||||
Содержимое:
|
||||
- Импорты всех необходимых модулей (модель, данные, утилиты, конфигурации).
|
||||
- Чтение гиперпараметров из config.yaml.
|
||||
- Основной цикл обучения.
|
||||
- Тестирование модели на валидационных данных.
|
||||
|
||||
## model.py — Определение архитектуры модели
|
||||
Этот файл содержит только архитектуру модели. Если необходимо сменить модель, достаточно будет изменить только этот файл.
|
||||
|
||||
Содержимое:
|
||||
- Определение слоёв и функции forward для прохождения данных через модель.
|
||||
|
||||
## datasets.py — Работа с данными
|
||||
Этот файл отвечает за загрузку, подготовку и трансформацию данных.
|
||||
|
||||
Содержимое:
|
||||
- Загрузка тренировочных и тестовых наборов данных.
|
||||
- Определение необходимых трансформаций.
|
||||
|
||||
## utils.py — Вспомогательные функции
|
||||
Этот файл содержит функции, которые помогают сократить основной код и избежать дублирования.
|
||||
|
||||
Содержимое:
|
||||
- Сохранение/загрузка модели.
|
||||
- Подсчёт точности.
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
batch_size: 128
|
||||
epochs: 10
|
||||
learning_rate: 0.001
|
||||
data_dir: "./data"
|
||||
logs_dir: "./logs"
|
||||
+19
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -0,0 +1,5 @@
|
||||
torch==2.0.1
|
||||
torchvision==0.15.2
|
||||
matplotlib==3.7.2
|
||||
pyyaml==6.0
|
||||
numpy==1.24.1
|
||||
@@ -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")
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user