利用pytorch实现图像识别demo

2023-03-26 11:54:40 浏览数 (1)

1.下载和安装PyTorch,以及所需的其他依赖项。

2.准备数据集,并将其转换为适合PyTorch使用的格式(例如,利用 torchvision 库中的 transform 处理图像数据,并将其转换为 tensor)。

3.创建一个神经网络模型,可以使用 torch.nn 模块中提供的各种层构建模型。

4.定义损失函数(如交叉熵损失函数)和优化器(如随机梯度下降优化器或 Adam 优化器)。

5.开始训练模型,将数据集分成训练集和验证集。在每个 epoch 中,使用训练集进行训练,并使用验证集进行模型性能评估。

6.使用训练好的模型进行预测,可以用模型对新的图片进行分类。

下面是一个简单的示例代码,它使用CNN模型对图像数据集进行分类

代码语言:javascript复制
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 定义神经网络模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv_layer = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(16, 32, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )

        self.fc_layer = nn.Sequential(
            nn.Linear(32*6*6, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(1024, 10),
        )

    def forward(self, x):
        x = self.conv_layer(x)
        x = x.view(-1, 32*6*6)
        x = self.fc_layer(x)
        return x

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
    transforms.RandomRotation(degrees=15),
    transforms.ColorJitter(),
    transforms.RandomHorizontalFlip(),
    transforms.CenterCrop(size=224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = datasets.ImageFolder(root='path/to/train/dataset', transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)

val_dataset = datasets.ImageFolder(root='path/to/val/dataset', transform=transform)
val_loader = DataLoader(dataset=val_dataset, batch_size=64, shuffle=False)

# 定义模型、损失函数和优化器
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 在每个 epoch 结束后进行模型性能评估
    with torch.no_grad():
        total = 0
        correct = 0
        for images, labels in val_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total  = labels.size(0)
            correct  = (predicted == labels).sum().item()

        accuracy = 100.0 * correct / total

        print(f'Epoch [{epoch   1}/{num_epochs}], '
              f'Loss: {loss.item():.4f}, '
              f'Accuracy: {accuracy:.2f}%')

# 使用训练好的模型进行预测
model.eval()
with torch.no_grad():
    outputs = model(new_image_tensor)
    _, predicted = torch.max(outputs.data, 1)

0 人点赞