Проблемы обучения нейронной сети при распознавании шахматных фигур ⇐ Python

[Anonymous]

Мне нужен совет относительно моего любимого проекта. Я создаю механизм распознавания шахматных фигур, используя библиотеку Python PyTorch.
Теперь у меня есть две основные проблемы:

• Мой набор данных.
  Как мне подготовить изображения?
  У меня 7 классов: король, конь, пешка, ладья, ферзь и пустая.
  Если я использую «torchvision» .datasets.ImageFolder» — как лучше всего распределить мои изображения по папке?
  Играет ли стиль изображения роль, например, если у меня есть графическая иллюстрация шахматной фигуры или изображение, сделанное с шахматной доски кусок - следует ли их разделить на две разные категории/классы?
  Кроме того, какие еще характеристики изображения играют роль при классификации изображений (цвет, градиенты и т. д.)?
  
  * также , количество изображений в каждой папке должно быть одинаковым?
• Мое обучение дает точность 94%. Однако мой прогноз дает неверные результаты.
  Мне интересно, есть ли помимо очевидного - настройка аугментации и обработка перегрузки; что еще можно сделать? Имеет ли это смысл и можно ли обрабатывать каждую категорию/класс отдельно (преобразования, слои)?

Вот мой код:
1. Само обучение находится в "neural_net.py":

Код: Выделить всё

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
from torch.utils.data import DataLoader
from PIL import Image
import random
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import seaborn as sns
class NeuralNet(nn.Module):
def __init__(self):
super(NeuralNet, self).__init__()
# Layer 1
self.conv1 = nn.Sequential(nn.Conv2d(3, 32,3),
nn.ReLU(),
nn.MaxPool2d(3, 2))
# Layer 2
self.conv2d_2 = nn.Conv2d(32, 64, (3,3))
self.relu_2 = nn.ReLU()
self.maxpool_2 = nn.MaxPool2d(3, 2)
# Layer 3
self.fc1 = nn.Linear(179776, 32)
self.fc2 = nn.Linear(32, 7)
def forward(self, x):
# Layer 1
y = self.conv1(x)
# Layer 2
y = self.conv2d_2(y)
y = self.relu_2(y)
y = self.maxpool_2(y)
# Layer 3
y = y.view(y.size(0), -1)
y = self.fc1(y)
y = self.fc2(y)
return y
def plot_confusion_matrix(cm, class_names):
plt.figure(figsize=(10, 7))
sns.heatmap(cm, annot=True, fmt='d', cmap="Blues", xticklabels=class_names, yticklabels=class_names)
plt.xlabel('Predicted')
plt.ylabel('True')
plt.title('Confusion Matrix')
plt.show()
def main():
print("Main block is running")
# Early stopping parameters
early_stopping_patience = 3
best_val_loss = float('inf')
epochs_no_improve = 0
batch_size = 8 #64
step_size = 7 #10
epoch_max = 40;
# Define transformations with various augmentations
transform = transforms.Compose([
transforms.Resize([224,224]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
from_google_drive = False
#google drive is accesible only on dev mode, not locally
if from_google_drive:
from google.colab import drive
drive.mount('/content/drive')
# Set the path to your dataset in Google Drive
train_data_path = '/content/drive/My Drive/dataset/train'
test_data_path = '/content/drive/My Drive/dataset/test'
# for local scan
else:
train_data_path = '/Users/sabina.livny/Pictures/dataset'
test_data_path = '/Users/sabina.livny/Pictures/dataset'
train_data = torchvision.datasets.ImageFolder(root=train_data_path, transform=transform)
test_data = torchvision.datasets.ImageFolder(root=test_data_path, transform=transform)
print("Class names:", train_data.classes)  # Check that classes are defined
# DataLoader for GPU
train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=True, num_workers=4)
# Device setup
if torch.backends.mps.is_available():  # For Apple Silicon GPUs (macOS)
device = torch.device("mps")
elif torch.cuda.is_available():  # For NVIDIA GPUs on Linux/Windows
device = torch.device("cuda")
else:
device = torch.device("cpu")  # Default to CPU
# Move the model to the selected device
net = NeuralNet().to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.1)  # Adjusted step_size
manual_train = True  # Set this to True to trigger training
if manual_train:
for epoch in range(epoch_max):
print(f'Training epoch {epoch}....')
net.train()
running_loss = 0.0
for i, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
avg_train_loss = running_loss / len(train_loader)
print(f'Training Loss:  {avg_train_loss:.4f}')
# Validation phase
net.eval()
val_loss = 0.0
with torch.no_grad():
for data in test_loader:
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
outputs = net(inputs)
loss = loss_function(outputs, labels)
val_loss += loss.item()
avg_val_loss = val_loss / len(test_loader)
print(f'Validation Loss: {avg_val_loss:.4f}')
# Early stopping check
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
epochs_no_improve = 0
torch.save(net.state_dict(), 'best_model.pth')
else:
epochs_no_improve += 1
if epochs_no_improve == early_stopping_patience:
print(f'Early stopping triggered at epoch {epoch}.')
break
# Save model every epoch (optional)
torch.save(net.state_dict(), 'trained_net.pth')
# Learning rate scheduler step
scheduler.step()
else:
print("Training skipped. Set manual_train to True to start training.")
# Manual flag for evaluating the model
evaluate_model = True  # Set to False to skip evaluation
compute_confusion_matrix = True
# List of class names
class_names = train_data.classes
if evaluate_model:
net.eval()
correct = 0
total = 0
all_labels = []
all_predictions = []
with torch.no_grad():
for data in test_loader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
if compute_confusion_matrix:
all_labels.extend(labels.cpu().numpy())
all_predictions.extend(predicted.cpu().numpy())
accuracy = 100 * correct / total
print(f'Accuracy: {accuracy:.2f}%')
if compute_confusion_matrix:
# Compute confusion matrix
cm = confusion_matrix(all_labels, all_predictions)
plot_confusion_matrix(cm, class_names)
if __name__ == '__main__':
main()

2. А это файл scan.py, в котором я делаю прогнозы:

Код: Выделить всё

# Make a prediction using the model
def predict(images):
transform = transforms.Compose([
transforms.Resize([224,224]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
predicted_classes = []
model = NeuralNet()
model.load_state_dict(torch.load('trained_net.pth'))  # Load the saved weights
model.eval()  # Set model to evaluation mode
with torch.no_grad():  # Disable gradient calculation
for image_path in images:
img = Image.open(image_path).convert('RGB')
img = transform(img)
img = img.unsqueeze(0)
img = img.to(next(model.parameters()).device)
output = model(img)  # Forward pass
_, predicted_class = torch.max(output, 1)  # Get the predicted class
predicted_classes.append(predicted_class.item())
return predicted_classes  # Return list of predicted class indices
def scan_pieces(file):
# Using the factory to create a piece
piece = PieceFactory.create_piece("rook")
piece.file = file  # This will trigger the setter
squares_images = get_images_cropped(file)
# You can call this to predict the class based on the processed image
predicted_classes = predict(squares_images)  # Make the prediction
# Print the class name
class_names = ['queen', 'rook', 'bishop', 'knight', 'pawn', 'king', 'empty']
# Print the predicted classes for all images
for index, predicted_class in enumerate(predicted_classes):
# Ensure the predicted_class index is within the range of class_names
if predicted_class < len(class_names):
print(f"Predicted Class for square {index + 1}: {class_names[predicted_class]}")
else:
print(f"Predicted Class for square {index + 1}: Unknown class (index {predicted_class})")

По сути, я хочу научить модель распознавать шахматные фигуры, чтобы позже можно было использовать эти веса для сравнения с загруженным изображением, взятым со снимка экрана некоторой шахматной позиции, который позже Я могу расположить фигуры на доске на основе загруженного изображения.
Вот моя матрица путаницы:




Подробнее здесь: https://stackoverflow.com/questions/791 ... ess-pieces