PyTorch框架的深度学习模型训练图像分类任务

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision.models as models
#1. 导入所需的库，如torch, torchvision等库，它们是Pytorch框架中的基本库，可以用来搭建神经网络，进行数据预处理等工作。 

# 定义网络
class Network(nn.Module):
#2. 定义一个 Network 类，继承自 nn.Module ，是自定义神经网络的基本结构。
    def __init__(self):
        super(Network, self).__init__()
#3. 在 Network 类中，初始化方法 __init__ 中，首先通过 super(Network, self).__init__() 调用父类的初始化方法。然后加载预训练的ResNet18模型。将预训练模型的参数设置为不可更新，以便在训练过程中保持不变。接着，修改模型的最后一层全连接层，将其输出特征数目从1000改为3，以适应我们要解决的问题。 
        # 加载预训练模型
        self.model = models.resnet18(pretrained=True)

        # 冻结前面的层
        for param in self.model.parameters():
            param.requires_grad = False

            # 修改最后一层
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, 3)

    def forward(self, x):
        x = self.model(x)
        return x
#4. 在 Network 类中，定义 forward 方法，用于前向传播，输入数据经过预训练的ResNet18模型处理后输出。

net = Network()
#5. 实例化 Network 对象，即创建一个神经网络对象。 
# 数据加载
train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(128),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0),
    transforms.RandomRotation(degrees=30),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(128),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
#6. 定义训练数据和验证数据的预处理操作，例如随机裁剪、水平翻转、归一化等。 
train_data = ImageFolder('datasets/train', transform=train_transforms)
val_data = ImageFolder('datasets/val', transform=val_transforms)
#7. 使用 ImageFolder 类创建训练数据和验证数据的数据集对象，并使用 transforms 参数应用预处理操作。 
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_data, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
#8. 使用 DataLoader 类创建训练数据和验证数据的数据加载器对象，用于将数据分批次加载进神经网络。 
# 损失函数
loss_func = nn.CrossEntropyLoss()
#9. 定义损失函数为交叉熵损失，用于评估网络的预测错误。 
# 优化器
optimizer = optim.Adam(net.parameters(), lr=0.001)
#10. 定义优化器为Adam，用于更新网络参数，优化损失函数。 
# 学习率控制进程
train_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
#11. 定义学习率调度器，用于在训练过程中按照策略调整学习率
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#12. 设置设备，如果有GPU则使用GPU，否则使用CPU。 
net = net.to(device)
#13. 将神经网络移动到指定设备上。 

def train(net, loss_func, optimizer, train_lr_scheduler, num_epochs=25):
#14. 定义 train 函数，用于训练和验证神经网络。训练过程中的每一轮（epoch）分为训练阶段和验证阶段，分别计算每个阶段的损失和准确率。在训练阶段，根据损失值更新神经网络的参数；在验证阶段，不更新参数，只计算损失和准确率。
    for epoch in range(num_epochs):
        print("Training:Epoch {}/{}".format(epoch, num_epochs - 1))
        net.train()
        epoch_loss = 0
        epoch_corrects = 0
        for inputs, labels in train_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = net(inputs)
            _, preds = torch.max(outputs, 1)
            loss = loss_func(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss = epoch_loss + loss.item() * inputs.size(0)
            epoch_corrects = epoch_corrects + torch.sum(preds == labels.data)
        epoch_loss = epoch_loss / len(train_data)
        epoch_corrects = epoch_corrects.double() / len(train_data)
        print('train loss: {}, acc: {}'.format(epoch_loss, epoch_corrects))
        train_lr_scheduler.step()
        print("Valuating:Epoch {}/{}".format(epoch, num_epochs - 1))
        net.eval()
        epoch_loss = 0
        epoch_corrects = 0
        for inputs, labels in val_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = net(inputs)
            _, preds = torch.max(outputs, 1)
            loss = loss_func(outputs, labels)
            epoch_loss = epoch_loss + loss.item() * inputs.size(0)
            epoch_corrects = epoch_corrects + torch.sum(preds == labels.data)
        epoch_loss = epoch_loss / len(val_data)
        epoch_corrects = epoch_corrects.double() / len(val_data)
        print('val loss: {}, acc: {}'.format(epoch_loss, epoch_corrects))

if __name__ == '__main__':
#15. 调用 train 函数，开始训练神经网络。
    train(net, loss_func, optimizer, train_lr_scheduler, num_epochs=25)