手把手教你用ResNet训练自定义数据集：含PNG标签、NPY图像，附正反向传播代码

使用ResNet训练自定义数据集的过程可以分为几个主要步骤，包括数据加载、模型定义、训练过程中的前向和反向传播。这是一个使用PyTorch实现的基本示例： 数据加载 首先，我们需要加载自定义数据集。假设你的数据集已经整理好，并创建了一个数据加载器。 ```python import torch import numpy as np import os from torchvision import transforms from torch.utils.data import Dataset, DataLoader from PIL import Image class CustomDataset(Dataset): def __init__(self, image_dir, label_dir, transform=None): self.image_dir = image_dir self.label_dir = label_dir self.transform = transform self.image_filenames = os.listdir(image_dir) def __len__(self): return len(self.image_filenames) def __getitem__(self, idx): img_name = os.path.join(self.image_dir, self.image_filenames[idx]) label_name = os.path.join(self.label_dir, self.image_filenames[idx].replace('.png', '.npy')) image = np.load(img_name) label = np.load(label_name) image = Image.fromarray(image) if self.transform: image = self.transform(image) return image, label # Dataset and DataLoader transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), ]) train_dataset = CustomDataset('path/to/train/images', 'path/to/train/labels', transform=transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) test_dataset = CustomDataset('path/to/test/images', 'path/to/test/labels', transform=transform) test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False) ``` 模型定义 接下来，我们定义ResNet模型。这里使用PyTorch提供的ResNet18作为基础。 ```python import torchvision.models as models import torch.nn as nn # 使用预训练的ResNet18模型 class ResNetModel(nn.Module): def __init__(self, num_classes): super(ResNetModel, self).__init__() self.model = models.resnet18(pretrained=True) self.model.fc = nn.Linear(self.model.fc.in_features, num_classes) def forward(self, x): return self.model(x) # 假设我们有10个类别 num_classes = 10 model = ResNetModel(num_classes) ``` 训练过程中的正向和反向传播 在训练过程中，我们将执行前向和反向传播。下面是训练的循环代码示例。 ```python import torch.optim as optim # 设置设备 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() # 假设进行多类别分类 optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练循环 num_epochs = 10 for epoch in range(num_epochs): model.train() # 设置模型为训练模式 running_loss = 0.0 for inputs, labels in train_loader: 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}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}") ``` 测试过程 测试过程通常不需要反向传播和参数更新。 ```python model.eval() # 设置模型为评估模式 with torch.no_grad(): # 关闭梯度计算 total_correct = 0 total_images = 0 for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) # 正向传播 _, predicted = torch.max(outputs.data, 1) # 获得预测结果 total_images += labels.size(0) total_correct += (predicted == labels).sum().item() print(f"Test Accuracy: {100 * total_correct / total_images:.2f}%") ``` 注意事项 1. 请根据自己的数据集类别数调整 `num_classes`。 2. 图像和标签的加载方式要根据具体数据格式进行调整，确保数据预处理的正确性。 3. 可以根据具体任务选择适合的损失函数。 4. 调整学习率、batch size、训练轮数等超参数以获得更好的训练效果。 5. 适时保存模型，以便后续加载和使用。 这是一个基本的实现示例，根据你的具体需求，可以进一步修改和优化。 ######[AI写代码神器 | 2766点数解答 | 2024-08-09 13:52:50]