一个Pytorch 的简单的分类本地图片的训练AI例子

下面是一个简单的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模型是一个简单的两层卷积层和一个全连接层的结构，适用于小规模的图像分类任务。

在实际应用中，你可能需要调整模型的结构、超参数以及使用更大的数据集进行训练以获得更好的性能。