Pytorch CNN简单实现 猫狗分类

---

CNN 猫狗分类

---

一、数据准备

准备要分类的图像数据：如猫狗数据（数据集）

二、数据载入

1.引入库

import os
import torch
import torch.nn as nn
from torch import optim
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
torch.manual_seed(1)  # 为了每次的实验结果一致

2.加载数据

class MyDataset(Dataset):
    def __init__(self, data, transform):
        self.data = data
        self.transform = transform

    def __getitem__(self, item):
        img, label = self.data[item]
        img = Image.open(img).convert('RGB')
        img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.data)


def load_data():
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.3),
        transforms.RandomVerticalFlip(p=0.3),
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))  # normalization
    ])
    #数据集存放路径
    cat_img_path = 'data/cat_dog/cat'
    dog_img_path = 'data/cat_dog/dog'
    train_img_list = [[os.path.join(cat_img_path, i), 0] for i in os.listdir(cat_img_path)[:160]] + [
        [os.path.join(dog_img_path, i), 1] for i in os.listdir(dog_img_path)[:160]]
    test_img_list = [[os.path.join(cat_img_path, i), 0] for i in os.listdir(cat_img_path)[160:]] + [
        [os.path.join(dog_img_path, i), 1] for i in os.listdir(dog_img_path)[160:]]
    train = MyDataset(train_img_list, transform=transform)
    test = MyDataset(test_img_list, transform=transform)
    train_data = DataLoader(dataset=train, batch_size=10, shuffle=True, num_workers=0)
    test_data = DataLoader(dataset=test, batch_size=1, shuffle=True, num_workers=0)
    return train_data, test_data

---

三、定义CNN模型

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        # 第一层卷积
        self.conv1 = nn.Sequential(
            # nn.Sequential:将有序的模块写在一起，将神经网络的层按顺序放在一起
            nn.Conv2d(
                in_channels=3,  # 输入图片的高度
                out_channels=16,  # 输出图片的高度
                kernel_size=3,  # 3x3的卷积核，相当于过滤器
                stride=2,  # 卷积核在图上滑动，每隔一个扫一次
                padding=1,  # 给图外边补上0
            ),
            # 数据归一化，使得数据在进行Relu之前不会因为数据过大而导致网络性能的不稳定
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        #
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                in_channels=16,
                out_channels=32,
                kernel_size=3,
                stride=2,
            ),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        #
        self.conv3 = nn.Sequential(
            nn.Conv2d(
                in_channels=32,
                out_channels=64,
                kernel_size=3,
                stride=2,
            ),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        # 输出层
        self.fc1 = nn.Linear(3 * 3 * 64, 64)
        self.fc2 = nn.Linear(64, 10)
        self.out = nn.Linear(10, 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        # print(x.size())
        x = x.view(x.shape[0], -1)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.out(x)
        return x

四、训练

def train():
    # 载入数据
    train_loader, test_loader = load_data()
    # 设置训练次数
    epoch_num = 20
    # 设置训练模式
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # cnn 实例化
    model = CNN().to(device)
    # 定义优化器和损失函数
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss().to(device)
    for epoch in range(epoch_num):
        for batch_idx, (data, target) in enumerate(train_loader, 0):
            data, target = Variable(data).to(device), Variable(target.long()).to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 10 == 0:
                print(f'Epoch:{epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                      f'({100. * batch_idx / len(train_loader):.2f}%)]\tLoss:{loss.item():.5f}')
    # 保存模型
    torch.save(model.state_dict(), "model/cnn_cat_dog.pkl")

五、测试

def test():
    train_loader, test_loader = load_data()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CNN().to(device)
    model.load_state_dict(torch.load("model/cnn_cat_dog.pkl"), False)
    model.eval()
    total = 0
    current = 0
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        predicted = torch.max(outputs.data, 1)[1].data
        total += labels.size(0)
        current += (predicted == labels).sum()
    print('Accuracy:%d%%' % (100 * current / total))

运行

if __name__ == '__main__':
    train() #训练
    # test() #测试