设置设备
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
定义数据增强
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
加载数据集
data_path = “path/to/dataset”
dataset = datasets.ImageFolder(root=data_path, transform=transform)
加载模型
model_path = “path/to/model.pth”
model = torch.load(model_path)
model.to(device)
model.eval()
获取预测结果和标签
labels = []
preds = []
for inputs, targets in dataset:
inputs = inputs.unsqueeze(0).to(device)
targets = targets.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
labels.append(targets.item())<