pytorch 实现图像识别

1. 代码讲解

1.1 导库

import os.path
from os import listdir
import numpy as np
import pandas as pd
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn import AdaptiveAvgPool2d
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import Dataset
import torchvision.transforms as transforms
from sklearn.model_selection import train_test_split

1.2 标准化、transform、设置GPU

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
normalize = transforms.Normalize(
   mean=[0.485, 0.456, 0.406],
   std=[0.229, 0.224, 0.225]
)
transform = transforms.Compose([transforms.ToTensor(), normalize])  # 转换

1.3 预处理数据

class DogDataset(Dataset):
# 定义变量
    def __init__(self, img_paths, img_labels, size_of_images):  
        self.img_paths = img_paths
        self.img_labels = img_labels
        self.size_of_images = size_of_images

# 多少长图片
    def __len__(self):
        return len(self.img_paths)

# 打开每组图片并处理每张图片
    def __getitem__(self, index):
        PIL_IMAGE = Image.open(self.img_paths[index]).resize(self.size_of_images)
        TENSOR_IMAGE = transform(PIL_IMAGE)
        label = self.img_labels[index]
        return TENSOR_IMAGE, label


print(len(listdir(r'C:\Users\AIAXIT\Desktop\DeepLearningProject\Deep_Learning_Data\dog-breed-identification\train')))
print(len(pd.read_csv(r'C:\Users\AIAXIT\Desktop\DeepLearningProject\Deep_Learning_Data\dog-breed-identification\labels.csv')))
print(len(listdir(r'C:\Users\AIAXIT\Desktop\DeepLearningProject\Deep_Learning_Data\dog-breed-identification\test')))

train_paths = []
test_paths = []
labels = []
# 训练集图片路径
train_paths_lir = r'C:\Users\AIAXIT\Desktop\DeepLearningProject\Deep_Learning_Data\dog-breed-identification\train'
for path in listdir(train_paths_lir):
    train_paths.append(os.path.join(train_paths_lir, path))  
# 测试集图片路径
labels_data = pd.read_csv(r'C:\Users\AIAXIT\Desktop\DeepLearningProject\Deep_Learning_Data\dog-breed-identification\labels.csv')
labels_data = pd.DataFrame(labels_data)  
# 把字符标签离散化，因为数据有120种狗，不离散化后面把数据给模型时会报错：字符标签过多。把字符标签从0-119编号
size_mapping = {}
value = 0
size_mapping = dict(labels_data['breed'].value_counts())
for kay in size_mapping:
    size_mapping[kay] = value
    value += 1
# print(size_mapping)
labels = labels_data['breed'].map(size_mapping)
labels = list(labels)
# print(labels)
print(len(labels))
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(train_paths, labels, test_size=0.2)

train_set = DogDataset(X_train, y_train, (32, 32))
test_set = DogDataset(X_test, y_test, (32, 32))

train_loader = torch.utils.data.DataLoader(train_set, batch_size=64)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=64)

1.4 建立模型

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()

        self.features = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5),  
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 120)
        )

    def forward(self, x):
        batch_size = x.shape[0]
        x = self.features(x)
        x = x.view(batch_size, -1)
        x = self.classifier(x)
        return x


model = LeNet().to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters())
TRAIN_LOSS = []  # 损失
TRAIN_ACCURACY = []  # 准确率

1.5 训练模型

def train(epoch):
    model.train()
    epoch_loss = 0.0 # 损失
    correct = 0  # 精确率
    for batch_index, (Data, Label) in enumerate(train_loader):
    # 扔到GPU中
        Data = Data.to(device)
        Label = Label.to(device)
        output_train = model(Data)
    # 计算损失
        loss_train = criterion(output_train, Label)
        epoch_loss = epoch_loss + loss_train.item()
    # 计算精确率
        pred = torch.max(output_train, 1)[1]
        train_correct = (pred == Label).sum()
        correct = correct + train_correct.item()
    # 梯度归零、反向传播、更新参数
        optimizer.zero_grad()
        loss_train.backward()
        optimizer.step()
    print('Epoch: ', epoch, 'Train_loss: ', epoch_loss / len(train_set), 'Train correct: ', correct / len(train_set))

1.6 测试模型

和训练集差不多。

def test():
    model.eval()
    correct = 0.0
    test_loss = 0.0
    with torch.no_grad():
        for Data, Label in test_loader:
            Data = Data.to(device)
            Label = Label.to(device)
            test_output = model(Data)
            loss = criterion(test_output, Label)
            pred = torch.max(test_output, 1)[1]
            test_correct = (pred == Label).sum()
            correct = correct + test_correct.item()
            test_loss = test_loss + loss.item()
    print('Test_loss: ', test_loss / len(test_set), 'Test correct: ', correct / len(test_set))

1.7结果

epoch = 10
for n_epoch in range(epoch):
    train(n_epoch)
test()

关于pytorch实现图像识别(实战)的文章就介绍至此，更多相关pytorch实现图像识别内容请搜索编程宝库以前的文章，希望以后支持编程宝库！

因为一些懂得都懂的原因，我的父母对于我电脑上的文件内容特别感兴趣。为了防止信息泄露，我连夜用Python做了一个文件加密器，防止我的重要信息被泄露。 下载地址github：https://gith ...