Pytorch分类模型的训练框架

PhotoDataset数据集是自己定义的数据集，数据集存放方式为：

----image文件夹

--------0文件夹

--------------img1.jpg

--------------img2.jpg

--------1文件夹

--------------img1.jpg

--------------img2.jpg

....

如果是cpu训练的话，就把代码中的.cuda()改成.cpu()

import os
import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from torch import nn
from torch import optim
from torch.utils.data.sampler import SubsetRandomSampler
import torch.nn.functional as F
import torchvision.models as models


# 定义数据集类
class PhotoDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.files = self._load_files()

    def _load_files(self):
        files = []
        for label in range(10):
            label_dir = os.path.join(self.root_dir, str(label))
            for file_name in os.listdir(label_dir):
                file_path = os.path.join(label_dir, file_name)
                files.append((file_path, int(label)))
        return files

    def __len__(self):
        return len(self.files)

    def __getitem__(self, index):
        file_path, label = self.files[index]
        image = torchvision.io.read_image(file_path, torchvision.io.ImageReadMode.RGB)
        image = image.float()
        if self.transform:
            image = self.transform(image)
        return image, label

# 数据集路径
data_dir = './photo2'

# 数据集预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    #transforms.ToTensor(),
    #transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 创建数据集
dataset = PhotoDataset(data_dir, transform=transform)

# 划分训练集和验证集
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
print(len(train_dataset))
print(len(val_dataset))
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)


# 定义基于ResNet-18的分类器
class ResNet18Classifier(nn.Module):
    def __init__(self, num_classes=10):
        super(ResNet18Classifier, self).__init__()
        # 使用预训练的ResNet-18模型
        self.resnet = models.resnet18(pretrained=False)
        # 替换掉原有的fc层，以适应我们的分类任务
        self.resnet.fc = nn.Linear(self.resnet.fc.in_features, num_classes)

    def forward(self, x):
        # 直接使用ResNet的forward函数
        x = self.resnet(x)
        return x

# 自定义模型，搭积木一样
class LightweightClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super(LightweightClassifier, self).__init__()
        # 定义卷积层
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        
        # 定义池化层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 定义全连接层
        self.fc1 = nn.Linear(64 * 32 * 32, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        # 通过卷积层和池化层
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        
        # 展平特征图以输入到全连接层
        x = x.view(-1, 64 * 32 * 32)
        
        # 通过全连接层
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        
        return x


# 创建模型实例
model = ResNet18Classifier() #LightweightClassifier()
model.cuda()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.Adam(model.parameters(), lr=0.0001)

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images = images.cuda()
        labels = labels.cuda()
        optimizer.zero_grad()
        outputs = model(images)
        #outputs = F.sigmoid(outputs)#二分类的话可以再加个sigmoid()
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')

    # 验证模型
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.cuda()
            labels = labels.cuda()
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        print(f'Validation Accuracy: {100 * correct / total}%')

其他

其实在这个框架中还有许多要素可以加入，比如：

1. Loss收敛趋势可视化，用tensorboard；【todo】
2. 学习率的动态调整，用各种warnup策略；【todo】
3. 对数据集的预处理，用albumentations库；【todo】
4. 使用T-SNE对数据集分布进行可视化；【todo】
5. 使用CAM可视化对模型提取到的特征进行可视化；【todo】
6. 模型权重的自动化保存策略，例如保存前10次效果最优的模型；【todo】