pytorch_图片分类_猫狗猴

图片分类作为cv基础，可以当作基础案例进行练习，这个是我自己做的一个猫狗猴的一个多分类，作为初学者可以试试，玩一下

import cv2
import warnings
import torch
import torchvision
import torchvision.transforms as transforms
import os
import random
import numpy as np
warnings.filterwarnings("ignore")
def transform_img(img):
    # 将图片尺寸缩放道 224x224
    img = cv2.resize(img, (32, 32))
    # 读入的图像数据格式是[H, W, C]
    # 使用转置操作将其变成[C, H, W]
    img = np.transpose(img, (2, 0, 1))
    img = img.astype('float32')
    # 将数据范围调整到[-1.0, 1.0]之间
    img = img / 255.
    img = img * 2.0 - 1.0
    return img


# 定义训练集数据读取器
def data_loader(datadir, batch_size=10, mode='train'):
    # 将datadir目录下的文件列出来，每条文件都要读入
    filenames = os.listdir(datadir)

    def reader():
        if mode == 'train':
            # 训练时随机打乱数据顺序
            random.shuffle(filenames)
        batch_imgs = []
        batch_labels = []
        for name in filenames:
            filepath = os.path.join(datadir, name)
            img = cv2.imread(filepath)
            img = transform_img(img)
            if name.split('_')[0] == 'D':
                label = 0
            elif name.split('_')[0] == 'H':
                label = 1
            elif name.split('_')[0] == 'T':
                # P开头的是病理性近视，属于正样本，标签为1
                label = 2
            else:
                raise ('Not excepted file name')
            # 每读取一个样本的数据，就将其放入数据列表中
            batch_imgs.append(img)
            batch_labels.append(label)
            if len(batch_imgs) == batch_size:
                # 当数据列表的长度等于batch_size的时候，
                # 把这些数据当作一个mini-batch，并作为数据生成器的一个输出
                imgs_array = np.array(batch_imgs).astype('float32')
                labels_array = np.array(batch_labels).astype('int64').reshape(-1, 1)
                yield imgs_array, labels_array
                batch_imgs = []
                batch_labels = []

        if len(batch_imgs) > 0:
            # 剩余样本数目不足一个batch_size的数据，一起打包成一个mini-batch
            imgs_array = np.array(batch_imgs).astype('float32')
            labels_array = np.array(batch_labels).astype('int64').reshape(-1, 1)
            yield imgs_array, labels_array

    return reader
DATADIR = 'data/dog_cat_pig'
train_loader = data_loader(DATADIR, batch_size=2, mode='train')
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 3)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
net = Net()
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for epoch in range(100):  # 多批次循环
    running_loss = 0.0
    for batch_id, data in enumerate(train_loader(),0):
        x_data, y_data = data
        x_data=torch.from_numpy(x_data)
        y_data=torch.from_numpy(y_data)
        inputs=x_data.to(device)
        labels=y_data.to(device).view(2)
        # 梯度置0
        optimizer.zero_grad()
        # 正向传播，反向传播，优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # 打印状态信息
        running_loss += loss.item()
    if epoch / 20 == 0:
        print(running_loss)
    torch.save(net.state_dict(), 'palm.pdparams')
print('Finished Training')
def predict(model, filepath, params_file_path):
    # 开启0号GPU预估
    use_gpu = False
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print('start evaluation .......')
    # 加载模型参数
    model_state_dict = torch.load(params_file_path)
    model.eval()
    batch_imgs = []
    img = cv2.imread(filepath)
    img = transform_img(img)
    imgs_array = np.array(img).astype('float32')
    batch_imgs.append(imgs_array)
    img = torch.tensor(batch_imgs)
    # 计算预测和精度
    result = model(img)
    pred = F.sigmoid(result)
    print(pred)
    dict_angle = {0: '狗', 1: '猴', 2: '猫'}
    for i, values in enumerate(pred[0]):
        if pred[0].max() == values:
            print(dict_angle[i])
    return result
predict(net, 'data/dog_cat_pig/T_4.jpg', params_file_path="palm.pdparams")

start evaluation …
猫
最后预测出来时一个猫，说明模型效果还可以，大家可以按照自己的想法，进行不同的分类，如有需要数据集和代码加群753035545