File size: 3,134 Bytes
411c80b 3b3bfbb 411c80b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
# Ustawienia parametrów treningu
img_width, img_height = 224, 224 # Wymiary obrazu wymagane przez model ResNet
batch_size = 32 # Liczba obrazów przetwarzanych na raz podczas treningu
epochs = 10 # Liczba epok treningu
learning_rate = 0.001 # Wskaźnik uczenia się dla optymalizatora
model_path = './result/animal_classifier_resnet.pth' # Ścieżka do zapisu wytrenowanego modelu
# Sprawdzenie, czy jest dostępny GPU i przypisanie urządzenia do zmiennej `device`
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Wyłączenie CuDNN, chyba że twoja karta wspiera bibliotekę CuDNN (NVIDIA CUDA Deep Neural Network library)
# torch.backends.cudnn.enabled = False
# Transformacje danych wejściowych (zmiana rozmiaru, normalizacja)
transform = transforms.Compose([
transforms.Resize((img_width, img_height)), # Zmiana rozmiaru obrazu
transforms.ToTensor(), # Konwersja obrazu do tensoru
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # Normalizacja obrazu
])
# Przygotowanie danych treningowych z katalogu `raw-img`
data_dir = 'raw-img' # Ścieżka do katalogu z obrazami treningowymi
train_dataset = datasets.ImageFolder(data_dir, transform=transform) # Wczytanie obrazów i zastosowanie transformacji
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # Tworzenie loadera danych
# Użycie pretrenowanego modelu ResNet18
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
num_ftrs = model.fc.in_features # Liczba wejściowych cech ostatniej warstwy
model.fc = nn.Linear(num_ftrs, len(train_dataset.classes)) # Zastąpienie ostatniej warstwy dopasowanej do liczby klas w danych
# Przeniesienie modelu na GPU, jeśli jest dostępny
model = model.to(device)
# Definicja funkcji kosztu (CrossEntropyLoss) i optymalizatora (Adam)
criterion = nn.CrossEntropyLoss() # Funkcja kosztu
optimizer = optim.Adam(model.parameters(), lr=learning_rate) # Optymalizator
# Trening modelu
for epoch in range(epochs): # Pętla przez epoki
model.train() # Ustawienie modelu w tryb treningowy
running_loss = 0.0 # Zmienna do śledzenia straty
for inputs, labels in train_loader: # Pętla przez batch'e danych
# Przeniesienie danych na GPU, jeśli jest dostępny
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad() # Zerowanie gradientów
outputs = model(inputs) # Przekazanie danych przez model
loss = criterion(outputs, labels) # Obliczenie straty
loss.backward() # Propagacja wsteczna
optimizer.step() # Aktualizacja wag modelu
running_loss += loss.item() # Akumulacja straty
print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader)}") # Wyświetlenie średniej straty na epokę
# Zapisywanie wytrenowanego modelu do pliku
torch.save(model.state_dict(), model_path)
print(f"Model saved to {model_path}")
|