下面是一个简单的PyTorch图像分类的例子，使用一个小型的卷积神经网络（CNN）对本地图片进行训练。在这个例子中，我们将使用torchvision来加载数据并定义一个简单的CNN模型进行训练。
首先，确保你已经安装了PyTorch和torchvision：
pip install torch torchvision
然后，使用以下代码创建一个简单的图像分类模型：
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义简单的CNN模型
class SimpleCNN(nn.Module):
def __init__(self, num_classes=2):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU()
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(32 * 64 * 64, num_classes) # 根据输入图像大小调整线性层输入大小
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = self.pool(x)
x = self.conv2(x)
x = self.relu(x)
x = self.pool(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
return x
# 定义数据转换
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
])
# 加载数据集
train_dataset = datasets.ImageFolder(root="path/to/train/dataset", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# 初始化模型、损失函数和优化器
model = SimpleCNN(num_classes=len(train_dataset.classes))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}")
# 保存模型
torch.save(model.state_dict(), "simple_cnn_model.pth")
请确保替换"path/to/train/dataset"为包含你训练数据的文件夹路径。这个例子中的CNN模型是一个简单的两层卷积层和一个全连接层的结构，适用于小规模的图像分类任务。
在实际应用中，你可能需要调整模型的结构、超参数以及使用更大的数据集进行训练以获得更好的性能。
