### PyTorch 示例代码：适用于 16GB 显存的 GPU 以下是一个完整的 PyTorch 实现，旨在展示如何在具有 16GB 显存的 GPU 上高效运行深度学习任务。此代码涵盖了模型定义、数据加载以及训练过程中的最佳实践。 ```python import torch from torch import nn, optim from torchvision import datasets, transforms from torch.utils.data import DataLoader # 检测是否有可用的 GPU 并设置默认设备 device = torch.device("cuda" if torch.cuda.is_available() and torch.cuda.get_device_properties(0).total_memory >= 1e10 else "cpu") # 数据预处理与增强 data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } # 下载 CIFAR-10 数据集并应用转换 dataset_train = datasets.CIFAR10(root='./data', train=True, download=True, transform=data_transforms['train']) dataset_val = datasets.CIFAR10(root='./data', train=False, download=True, transform=data_transforms['val']) dataloader_train = DataLoader(dataset_train, batch_size=64, shuffle=True, num_workers=4) dataloader_val = DataLoader(dataset_val, batch_size=64, shuffle=False, num_workers=4) # 定义卷积神经网络模型 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv_layers = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.fc_layers = nn.Sequential( nn.Linear(64 * 8 * 8, 512), nn.ReLU(), nn.Dropout(p=0.5), nn.Linear(512, 10) # CIFAR-10 类别数量为 10 ) def forward(self, x): x = self.conv_layers(x) x = torch.flatten(x, start_dim=1) x = self.fc_layers(x) return x model = SimpleCNN().to(device) # 设置损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练循环 def train_model(dataloaders, model, criterion, optimizer, epochs=10): for epoch in range(epochs): model.train() running_loss = 0.0 for inputs, labels in dataloaders["train"]: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(dataloaders["train"])}') # 验证模型性能 def validate_model(dataloader, model): model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in dataloader: images, labels = images.to(device), 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 print(f'Validation Accuracy: {accuracy}%') # 调用训练和验证功能 train_model({"train": dataloader_train}, model, criterion, optimizer, epochs=5) validate_model(dataloader_val, model) ``` 上述代码实现了以下几个关键点： 1. **设备检测**：自动判断当前环境中是否存在满足条件（至少 16GB 显存）的 GPU，并相应地分配计算资源[^1]。 2. **数据增强**：通过对图像进行随机裁剪、翻转等操作提高模型泛化能力。 3. **批量大小调整**：考虑到目标 GPU 的显存量较大，适当增加批次大小至 64 可有效减少每轮迭代次数，加快收敛速度。 4. **Dropout 层引入**：防止过拟合现象发生的同时保持较高的测试精度水平[^3]。 --- ###