### 使用 PyTorch 训练 ResNet18 模型的代码示例 为了展示如何使用 PyTorch 训练 ResNet18 模型，下面提供了一个完整的代码示例。此代码涵盖了数据加载、模型定义、训练循环以及评估过程。 #### 导入必要的库 ```python import torch from torchvision import models, datasets, transforms import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader ``` #### 定义设备配置 ```python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f'Using device: {device}') ``` #### 数据预处理与增强 ```python transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) ``` #### 加载 CIFAR-10 数据集 ```python trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) ``` #### 初始化 ResNet18 模型并调整最后分类层 ```python model = models.resnet18(pretrained=False) num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10) # 修改最后一层以适应 CIFAR-10 类别数 model.to(device) ``` #### 设置损失函数和优化器 ```python criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) ``` #### 训练模型 ```python def train_model(model, criterion, optimizer, scheduler, num_epochs=25): best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs -' * 10) model.train() running_loss = 0.0 running_corrects = 0 for inputs, labels in trainloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) scheduler.step() epoch_loss = running_loss / len(trainset) epoch_acc = running_corrects.double() / len(trainset) print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') test_accuracy = evaluate_model(model, testloader) print(f'Test Accuracy: {test_accuracy:.4f}\n') if test_accuracy > best_acc: best_acc = test_accuracy best_model_wts = model.state_dict() model.load_state_dict(best_model_wts) return model def evaluate_model(model, dataloader): model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in dataloader: images = images.to(device) labels = labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total return accuracy ``` #### 开始训练 ```python if __name__ == '__main__': trained_model = train_model(model, criterion, optimizer, scheduler, num_epochs=25) ```