基于 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数据集,并进行训练和测试。请根据您的具体需求调整模型结构、训练轮数、学习率等参数。如果您的任务更加复杂或数据集不同,可能需要更复杂的模型架构和训练策略。

相关推荐
DataCastle3 分钟前
第三届Bio-OS AI开源大赛启动会隆重举行
人工智能
后端小肥肠12 分钟前
躺赚必备!RPA+Coze+豆包:公众号自动发文,AI率0%亲测有效(附AI率0%提示词)
人工智能·aigc·coze
摘星编程23 分钟前
CloudBase AI ToolKit实战:从0到1开发一个智能医疗网站
人工智能·腾讯云·ai代码远征季#h5应用·ai医疗应用·cloudbase开发
锅挤28 分钟前
深度学习5(深层神经网络 + 参数和超参数)
人工智能·深度学习·神经网络
一支烟一朵花32 分钟前
630,百度文心大模型4.5系列开源!真香
人工智能·百度·开源·文心一言
网安INF34 分钟前
深层神经网络:原理与传播机制详解
人工智能·深度学习·神经网络·机器学习
AIbase20241 小时前
国内MCP服务平台推荐!aibase.cn上线MCP服务器集合平台
运维·服务器·人工智能
喜欢吃豆2 小时前
快速手搓一个MCP服务指南(九): FastMCP 服务器组合技术:构建模块化AI应用的终极方案
服务器·人工智能·python·深度学习·大模型·github·fastmcp
星融元asterfusion2 小时前
基于路径质量的AI负载均衡异常路径检测与恢复策略
人工智能·负载均衡·异常路径