torchvision构建行人重识别模型机器视觉.zip

import torch import torch.nn as nn import torch.optim as optim import torchvision.models as models import torchvision.transforms as transforms from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader from torch.optim.lr_scheduler import StepLR # 数据增强 train_transform = transforms.Compose([ transforms.RandomResizedCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 加载训练集和测试集 train_dataset = ImageFolder(root='dataset/train', transform=train_transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_transform = transforms.Compose([ transforms.Resize((256, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_dataset = ImageFolder(root='dataset/test', transform=test_transform) test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False) # 定义行人重识别模型，使用ResNet-50作为特征提取模型 class PersonReIDModel(nn.Module): def __init__(self, num_classes): super(PersonReIDModel, self).__init__() self.resnet = models.resnet50(pretrained=True) in_features = self.resnet.fc.in_features self.resnet.fc = nn.Linear(in_features, num_classes) def forward(self, x): return self.resnet(x) # 初始化模型和优化器 model = PersonReIDModel(num_classes=len(train_dataset.classes)) optimizer = optim.Adam(model.parameters(), lr=0.001) scheduler = StepLR(optimizer, step_size=5, gamma=0.1) # 学习率调整策略 criterion = nn.CrossEntropyLoss() # 训练模型 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) epochs = 10 for epoch in range(epochs): model.train() running_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_loader)}") # 调整学习率 scheduler.step() print("Training finished!") # 保存模型 torch.save(model.state_dict(), 'person_reid_model.pth') # 在测试集上评估模型 model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in test_loader: images, labels = images