CNN分类网络之：Lenet

落日映苍穹つ 2023-02-14 13:04 5阅读 0赞

pytorch 写的一个 lenet的分类网络，不是百分百还原哈，结构是一样的， 简单训练一下自己的数据集。数据集格式如下，data内存放 自己的数据，每个类别放到一个文件夹中，文件夹名称为类别标签如下图

![20200602212822824.png][]

1.网络搭建

import torch
    import torch.nn as nn
    
    
    class Lenet(nn.Module):
        def __init__(self, num_classes = 1000):
            super(Lenet, self).__init__()
            self.conv1 = nn.Conv2d(3,6,5,1,0)
            self.pool2 = nn.MaxPool2d(2)
            self.conv3 = nn.Conv2d(6, 16, 5, 1, 0)
            self.pool4 = nn.MaxPool2d(2)
            self.conv5 = nn.Conv2d(16, 120, 5, 1, 0)
            self.fc6 = nn.Linear(120, 84)
            self.fc7 = nn.Linear(84, 10)
            self.relu = nn.ReLU(inplace=True)
    
        def forward(self, x):
            x = self.conv1(x)
            x = self.relu(x)
            x = self.pool2(x)
    
            x = self.conv3(x)
            x = self.relu(x)
            x = self.pool4(x)
    
            x = self.conv5(x)
            x = self.relu(x)
    
            x = x.view(x.size()[0], -1)
    
            x = self.fc6(x)
            x = self.relu(x)
            x = self.fc7(x)
            return x

2、读取自己的图像及标签数据

from torch.utils.data import Dataset
    import cv2
    import os
    from torchvision import transforms as tvtsf
    import torch
    import numpy as np
    
    class cnndata(Dataset):
        def __init__(self, path):
            self.path = path
            self.img, self.cls, self.class_label= self.get_filname_and_cls()
            print("图像共有：%s 张"%(len(self.img)))
    
        def __getitem__(self, item):
            img = cv2.imread(self.img[item], 1)
            img_src = cv2.resize(img, (32, 32), cv2.INTER_AREA)
    
            if img_src.shape[2] == 1:
                img_src = cv2.cvtColor(img_src, cv2.COLOR_GRAY2BGR)
            image = torch.from_numpy(img_src / 255.)
            image = image.permute(2, 0, 1).contiguous()
            label = torch.from_numpy(np.array(self.cls[item]))
    
            return image, label
    
        def get_filname_and_cls(self):
            if not os.path.exists(self.path):
                raise Exception("no wenjianjia")
    
            class_name = os.listdir(self.path)
    
            cls = []
            imgs = []
            class_label = {}
            for c, cl in enumerate(class_name):
                class_label[c] = cl
                filename = os.listdir(os.path.join(self.path, cl))
                filename.sort()
                for name in filename:
                    # img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
                    imgs.append(os.path.join(self.path, cl, name))
                    cls.append(c)
    
            return imgs, cls, class_label
    
        def get_classlabel(self):
            return self.class_label
    
        def __len__(self):
            return len(self.img)

3、 训练及测试

from get_data import cnndata
    from torch.utils.data import DataLoader
    import torch.nn as nn
    from torch import optim
    import torch
    import time
    from Lenet import Lenet
    
    
    def train(batch=16, traindata=None, model= None, epochs=20):
        dataloader = DataLoader(traindata, batch_size=int(batch), shuffle=True,num_workers=2)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters())
    
        for epoch in range(int(epochs)):
            star = time.time()
            correct_num = 0
            epoch_loss = 0
            for img, cls in dataloader:
                img= img.type(torch.FloatTensor).cuda()
                cls = cls.type(torch.LongTensor).cuda()
    
                optimizer.zero_grad()
    
                out = model(img).cuda()
                loss = criterion(out, cls)
    
                epoch_loss += loss.detach().cpu().numpy()
                correct_num += torch.eq(cls, out.argmax(dim=1)).sum().item()
    
                loss.backward()
                optimizer.step()
    
            print("Epoch :%d , loss%.4f, acc:%.3f, time:%.3f"
                  %(epoch, round(epoch_loss/len(traindata), 4), round(correct_num/(len(traindata)), 3), time.time()-star))
            if correct_num/len(traindata) > 0.8:
                torch.save(model.state_dict(), 'lenet.pth')
    
    
    def test(batch=16, traindata=None, model= None, class_label =None):
        dataloader = DataLoader(traindata, batch_size=int(batch), shuffle=True, num_workers=2)
        model.load_state_dict(torch.load('lenet.pth'))
        model.eval()
        with torch.no_grad():
            star = time.time()
            for img, cls in dataloader:
                img = img.type(torch.FloatTensor).cuda()
                cls = cls.type(torch.LongTensor).cuda()
                out = model(img).cuda()
                pre = out.argmax(dim=1)
                print('pre is:', pre.detach().cpu().numpy(),'label is:', cls.detach().cpu().numpy())
    
    
    def weights_init(m):
        if isinstance(m, nn.Conv2d):
            nn.init.xavier_uniform_(m.weight.data)
            m.bias.data.zero_()
            
            
    if __name__ == '__main__':
        batch = 2
        traindata = cnndata(path="./data")
    
        model = Lenet(num_classes=10).cuda()
        model.apply(weights_init)
        class_label =traindata.get_classlabel()
        print('index and labels:', class_label)
        train(batch=batch, traindata=traindata, model=model)
        # test(batch=batch, traindata=traindata, model=model, class_label=class_label)

训练结果：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg_size_16_color_FFFFFF_t_70][]

预测结果：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg_size_16_color_FFFFFF_t_70 1][]

[20200602212822824.png]: https://img-blog.csdnimg.cn/20200602212822824.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200602211905767.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg==,size_16,color_FFFFFF,t_70
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg_size_16_color_FFFFFF_t_70 1]: https://img-blog.csdnimg.cn/2020060221254316.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RZdGFuZ3lhbg==,size_16,color_FFFFFF,t_70