图像分类 pytorch

PyTorch是一个流行的深度学习框架，可以用来进行图像分类任务。下面是一个简单的图像分类示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, models, transforms

# 数据预处理和加载
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_data = datasets.ImageFolder('path_to_train_data', transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)

# 加载预训练模型
model = models.resnet50(pretrained=True)
num_classes = len(train_data.classes)
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
for epoch in range(10):
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = 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}: Loss {running_loss/len(train_loader)}")

# 在测试集上评估模型
test_data = datasets.ImageFolder('path_to_test_data', transform=transform)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=False)

correct = 0
total = 0
with torch.no_grad():
    model.eval()
    for inputs, labels in test_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy on test set: {correct/total}")

在上述示例代码中，我们首先对图像数据进行预处理和加载。然后使用预训练的ResNet-50模型进行迁移学习，将最后一层全连接层替换成适合分类任务的新层。然后定义损失函数和优化器。接下来，在训练集上进行模型训练，并在每个epoch打印损失值。最后，在测试集上评估模型的分类准确率。

请注意，上述代码中的`'path_to_train_data'`和`'path_to_test_data'`需要替换为您自己的数据集路径。此外，还可以根据需要调整超参数和其他设置。希望这个示例对您有帮助！