使用pytorch训练一个简单的LeNet5分类器

"""赞美欧姆尼塞娅"""

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import os

# LeNet-5 Architecture
class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 输入是三通道（RGB），所以将第一层卷积的输入通道改为 3
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)  # 输入: 3x32x32, 输出: 6x28x28
        self.pool = nn.MaxPool2d(2, 2)               # 池化大小 2x2
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)  # 输入: 6x14x14, 输出: 16x10x10
        self.fc1 = nn.Linear(16 * 5 * 5, 120)         # 全连接层
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 3)                  # 输出层，有 3 个类别

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # 卷积1 -> ReLU -> 池化
        x = self.pool(torch.relu(self.conv2(x)))  # 卷积2 -> ReLU -> 池化
        x = x.view(-1, 16 * 5 * 5)  # 将输出展平
        x = torch.relu(self.fc1(x))  # 全连接层 1
        x = torch.relu(self.fc2(x))  # 全连接层 2
        x = self.fc3(x)              # 输出层
        return x

# Evaluate function to calculate accuracy on the validation set
def evaluate_model(model, dataloader, device='cpu'):
    model.eval()  # Set the model to evaluation mode
    correct = 0
    total = 0
    with torch.no_grad():  # No need to track gradients during evaluation
        for inputs, labels in dataloader:
            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()
    
    accuracy = correct / total  # Calculate accuracy
    return accuracy

# Training function
def train_model(model, trainloader, valloader, criterion, optimizer, num_epochs, device='cpu'):
    model.to(device)
    
    best_val_accuracy = 0.0

    for epoch in range(num_epochs):
        model.train()  # Set the model to training mode
        running_loss = 0.0
        for inputs, labels in trainloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()  # Zero the gradients
            outputs = model(inputs)  # Forward pass
            loss = criterion(outputs, labels)  # Compute loss
            loss.backward()  # Backward pass
            optimizer.step()  # Update weights
            
            running_loss += loss.item()
        
        # Print training loss for each epoch
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(trainloader)}")
        
        # Validation after each epoch
        val_accuracy = evaluate_model(model, valloader, device)
        print(f"Epoch {epoch+1}/{num_epochs}, Validation Accuracy: {val_accuracy:.4f}")
        
        # Save the best model based on validation accuracy
        if val_accuracy > best_val_accuracy:
            best_val_accuracy = val_accuracy
            # Save the best model
            torch.save(model.state_dict(), './best_lenet5_model.pth')

    print("Training Finished!")
    return best_val_accuracy

# Main function to set up data, model, optimizer, and train
def main():
    # Check if GPU is available
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Data augmentation and normalization
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # Simple normalization
    ])
    
    # Data directories (ensure you have train and val folders)
    train_dir = './data/train'
    val_dir = './data/val'

    # Load datasets using ImageFolder
    trainset = torchvision.datasets.ImageFolder(root=train_dir, transform=transform)
    valset = torchvision.datasets.ImageFolder(root=val_dir, transform=transform)

    # Data loaders
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)
    valloader = torch.utils.data.DataLoader(valset, batch_size=32, shuffle=False)

    # Initialize model, loss function, and optimizer
    model = LeNet5()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)  # Use Adam optimizer
    
    # Train the model and save the best model
    best_accuracy = train_model(model, trainloader, valloader, criterion, optimizer, num_epochs=2000, device=device)
    
    # Load and test the best model
    model.load_state_dict(torch.load('./best_lenet5_model.pth'))
    final_accuracy = evaluate_model(model, valloader, device)
    print(f"Best model validation accuracy: {best_accuracy:.4f}")
    print(f"Final model validation accuracy: {final_accuracy:.4f}")

if __name__ == '__main__':
    main()

训练集等结构如下：

内部是分类图片，文件名=类名：

测试用的是尺寸比较小的缺陷图像，测试结果也还不错