CNN简单实战：pytorch搭建CNN对猫狗图片进行分类

缺乏、安全感 2022-11-29 13:21 263阅读 0赞

在上一篇文章：[CNN训练前的准备：pytorch处理自己的图像数据（Dataset和Dataloader）][CNN_pytorch_Dataset_Dataloader]，大致介绍了怎么利用pytorch把猫狗图片处理成CNN需要的数据，今天就用该数据对自己定义的CNN模型进行训练及测试。

*  首先导入需要的包：

import torch
    from torch import optim
    import torch.nn as nn
    from torch.autograd import Variable
    from torchvision import transforms
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image

*  定义自己的CNN网络

class cnn(nn.Module):
        def __init__(self):
            super(cnn, self).__init__()  # 继承__init__功能
            # 第一层卷积
            self.conv1 = nn.Sequential(
                nn.Conv2d(
                    in_channels=3,
                    out_channels=32,
                    kernel_size=3,
                    stride=1,
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            # 第二层卷积
            self.conv2 = nn.Sequential(
                nn.Conv2d(
                    in_channels=32,
                    out_channels=64,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            # 第三层卷积
            self.conv3 = nn.Sequential(
                nn.Conv2d(
                    in_channels=64,
                    out_channels=128,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
                # nn.Dropout(0.1)
            )
    
            # 第四层卷积
            self.conv4 = nn.Sequential(
                nn.Conv2d(
                    in_channels=128,
                    out_channels=256,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            self.output = nn.Linear(in_features=256 * 14 * 14, out_features=2)
    
        def forward(self, x):
            x = self.conv1(x)
            x = self.conv2(x)
            x = self.conv3(x)
            x = self.conv4(x)
            temp = x.view(x.shape[0], -1)
            output = self.output(temp)
            return output, x

*  训练（GPU）

def train():
        train_loader, test_loader = load_data()
        epoch_num = 3
        # GPU计算
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = cnn().to(device)
        optimizer = optim.Adam(model.parameters(), lr=0.00005)
        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()  # 梯度清0
                output = model(data)[0]  # 前向传播
                loss = criterion(output, target)  # 计算误差
                loss.backward()  # 反向传播
                optimizer.step()  # 更新参数
                if batch_idx % 10 == 0:
                    print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        epoch, batch_idx * len(data), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader), loss.item()))
    
        torch.save(model, 'cnn.pkl')

一共训练三轮，训练的步骤如下：

1.  初始化模型：

model = cnn().to(device)

1.  选择优化器以及优化算法，这里选择了Adam：

optimizer = optim.Adam(model.parameters(), lr=0.00005)

1.  选择损失函数，这里选择了交叉熵：

criterion = nn.CrossEntropyLoss().to(device)

1.  对每一个batch里的数据，先将它们转成能被GPU计算的类型：

data, target = Variable(data).to(device), Variable(target.long()).to(device)

1.  梯度清零、前向传播、计算误差、反向传播、更新参数：

optimizer.zero_grad()  # 梯度清0
    output = model(data)[0]  # 前向传播
    loss = criterion(output, target)  # 计算误差
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

*  测试（GPU）

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

由于种种原因（训练集过小，网络定义不合理，学习率设置不合理…），结果惨不忍睹：  
![在这里插入图片描述][20200821192654949.png_pic_center]  
还需改进！！  
附上完整代码：

# -*- coding: utf-8 -*-
    """ @Time ： 2020/8/18 9:11 @Author ：KI @File ：CNN.py @Motto：Hungry And Humble """
    import torch
    from torch import optim
    import torch.nn as nn
    from torch.autograd import Variable
    from torchvision import transforms
    from torch.utils.data import Dataset, DataLoader
    from PIL import Image
    
    
    def Myloader(path):
        return Image.open(path).convert('RGB')
    
    
    # 得到一个包含路径与标签的列表
    def init_process(path, lens):
        data = []
        name = find_label(path)
        for i in range(lens[0], lens[1]):
            data.append([path % i, name])
    
        return data
    
    
    class MyDataset(Dataset):
        def __init__(self, data, transform, loder):
            self.data = data
            self.transform = transform
            self.loader = loder
    
        def __getitem__(self, item):
            img, label = self.data[item]
            img = self.loader(img)
            img = self.transform(img)
            return img, label
    
        def __len__(self):
            return len(self.data)
    
    
    def find_label(str):
        first, last = 0, 0
        for i in range(len(str) - 1, -1, -1):
            if str[i] == '%' and str[i - 1] == '.':
                last = i - 1
            if (str[i] == 'c' or str[i] == 'd') and str[i - 1] == '/':
                first = i
                break
    
        name = str[first:last]
        if name == 'dog':
            return 1
        else:
            return 0
    
    
    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))  # 归一化
        ])
        path1 = 'cnn_data/data/training_data/cats/cat.%d.jpg'
        data1 = init_process(path1, [0, 500])
        path2 = 'cnn_data/data/training_data/dogs/dog.%d.jpg'
        data2 = init_process(path2, [0, 500])
        path3 = 'cnn_data/data/testing_data/cats/cat.%d.jpg'
        data3 = init_process(path3, [1000, 1200])
        path4 = 'cnn_data/data/testing_data/dogs/dog.%d.jpg'
        data4 = init_process(path4, [1000, 1200])
        # 1300个训练
        train_data = data1 + data2 + data3[0:150] + data4[0:150]
    
        train = MyDataset(train_data, transform=transform, loder=Myloader)
        # 100个测试
        test_data = data3[150:200] + data4[150:200]
        test = MyDataset(test_data, transform=transform, loder=Myloader)
    
        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
    
    
    class cnn(nn.Module):
        def __init__(self):
            super(cnn, self).__init__()  # 继承__init__功能
            # 第一层卷积
            self.conv1 = nn.Sequential(
                nn.Conv2d(
                    in_channels=3,
                    out_channels=32,
                    kernel_size=3,
                    stride=1,
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            # 第二层卷积
            self.conv2 = nn.Sequential(
                nn.Conv2d(
                    in_channels=32,
                    out_channels=64,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            # 第三层卷积
            self.conv3 = nn.Sequential(
                nn.Conv2d(
                    in_channels=64,
                    out_channels=128,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
    
            # 第四层卷积
            self.conv4 = nn.Sequential(
                nn.Conv2d(
                    in_channels=128,
                    out_channels=256,
                    kernel_size=3,
                    stride=1,
                    padding=0
                ),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2),
            )
            self.output = nn.Linear(in_features=256 * 14 * 14, out_features=2)
    
        def forward(self, x):
            x = self.conv1(x)
            x = self.conv2(x)
            x = self.conv3(x)
            x = self.conv4(x)
            temp = x.view(x.shape[0], -1)
            output = self.output(temp)
            return output, x
    
    
    def train():
        train_loader, test_loader = load_data()
        epoch_num = 3
        # GPU计算
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = cnn().to(device)
        optimizer = optim.Adam(model.parameters(), lr=0.00005)
        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()  # 梯度清0
                output = model(data)[0]  # 前向传播
                loss = criterion(output, target)  # 计算误差
                loss.backward()  # 反向传播
                optimizer.step()  # 更新参数
                if batch_idx % 10 == 0:
                    print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        epoch, batch_idx * len(data), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader), loss.item()))
    
        torch.save(model, 'cnn.pkl')
    
    
    def test():
        train_loader, test_loader = load_data()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = torch.load('cnn.pkl')  # load model
        total = 0
        current = 0
        for data in test_loader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)[0]
    
            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()

[CNN_pytorch_Dataset_Dataloader]: https://blog.csdn.net/cyril_ki/article/details/108098617
[20200821192654949.png_pic_center]: /images/20221124/dd3f47013e0a4e378c041f40f4233114.png