基于 PyTorch 框架写一个图片分类模型

根据您的需求,我将提供一个基于PyTorch框架的简单图片分类模型示例。这个例子使用的是卷积神经网络(CNN)来对CIFAR-10数据集进行分类,CIFAR-10包含60,000个32x32彩色图像,分为10个类别。请注意,实际应用中您可能需要根据具体任务调整模型结构、训练参数等。

首先,请确保您已经安装了PyTorch和torchvision库。如果没有安装,可以通过以下命令安装:

Shell 复制代码
pip install torch torchvision

下面是基本的模型定义和训练代码:

python 复制代码
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):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        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.bn2 = nn.BatchNorm2d(32)
        self.fc = nn.Linear(32 * 8 * 8, 10)  # 假设经过两次池化后尺寸变为8x8

    def forward(self, x):
        x = self.pool(self.relu(self.bn1(self.conv1(x))))
        x = self.pool(self.relu(self.bn2(self.conv2(x))))
        x = x.view(-1, 32 * 8 * 8)  # 扁平化
        x = self.fc(x)
        return x

# 数据预处理
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR-10数据集
trainset = datasets.CIFAR10(root='./data', train=True,
                            download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2)

testset = datasets.CIFAR10(root='./data', train=False,
                           download=True, transform=transform)
testloader = DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

# 初始化模型、损失函数和优化器
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), 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 / (i + 1)}')

print('Finished Training.')

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

这段代码展示了如何定义一个简单的CNN模型,加载CIFAR-10数据集,并进行训练和测试。请根据您的具体需求调整模型结构、训练轮数、学习率等参数。如果您的任务更加复杂或数据集不同,可能需要更复杂的模型架构和训练策略。

相关推荐
极限实验室7 分钟前
Coco AI 实战(二):摄入MongoDB 数据
人工智能·mongodb
AIGC包拥它21 分钟前
AI教学设计助手:生成好教案的Prompt技术实战(一)
人工智能·prompt
SoFlu软件机器人1 小时前
Cursor、飞算JavaAI、GitHub Copilot、Gemini CLI 等热门 AI 开发工具合集
人工智能·github·copilot
isNotNullX2 小时前
实时数仓和离线数仓还分不清楚?看完就懂了
大数据·数据库·数据仓库·人工智能·数据分析
Liudef062 小时前
大语言模型的极限:知识、推理与创造力的边界探析
人工智能·语言模型·自然语言处理
潮湿的心情2 小时前
亚洲牧原:活跃行业交流,延伸公益版图,市场拓展再结硕果
大数据·人工智能
平和男人杨争争2 小时前
机器学习14——线性回归
人工智能·机器学习·线性回归
一个天蝎座 白勺 程序猿2 小时前
飞算JavaAI进阶:重塑Java开发范式的AI革命
java·开发语言·人工智能
李昊哲小课2 小时前
pandas销售数据分析
人工智能·python·数据挖掘·数据分析·pandas
whabc1003 小时前
和鲸社区深度学习基础训练营2025年关卡2(2)sklearn中的MLPClassifier
人工智能·深度学习·numpy