Resnet34垃圾分类图像深度学习

发布于 2022-08-22  1321 次阅读


大一做的卷积神经网络的深度学习,代码是大概七月15日左右写完的,后续经过少许修改,先暂时做个总结。
任务是具体是实现四种垃圾的分类,采用torchvison中的resnet34的模型和models.ResNet34_Weights.IMAGENET1K_V1的初始化参数。主要技术方面在于数据预处理以便于提升速度和在我的RTX3060的笔记本上能够进行训练。
水平有限,只是写个玩玩,只有数据预处理的代码感觉还行(全部我自己写的部分),其他都或多或少参考了一下github和资料书的内容,代码的注释应该还算比较详细了,按道理非常好懂
暂时未完成训练 只是做个记录

模型和算法用的是现成的算法,torchvison库里面的resnet34。因此,核心代码:

    net = models.resnet34(weights = models.ResNet34_Weights.IMAGENET1K_V1)

训练准备概览

    # ./main.py
    # 采用resnet34网络
    # 核心代码!!!
    net = models.resnet34(weights = models.ResNet34_Weights.IMAGENET1K_V1)
    net.fc = nn.Sequential(
        nn.Linear(512, 4)
        # nn.ReLU(),
        # nn.Dropout(0.5),
        # nn.Linear(num_hiddens1, num_hiddens2),
        # nn.ReLU(),
        # nn.Dropout(0.5),
        # nn.Linear(num_hiddens2, 4)
        # nn.Linear(256, 4)
    )
    # net.load_state_dict(torch.load('./model_5.pt'))

最初打算采用的是resnet50,但是因为数据集的原因发现过拟合现象相当严重,因此换为这个模型,并且删除了dropout层。

训练的模型参数为:

# ./config.py
# 存储常量 变量和超参数
import torch

# 训练所需参数
batch_size = 96
num_epochs = 8000
num_workers = 1
lr_init = 0.001
lr_stepsize = 20
weight_decay = 0.01
num_hiddens1 = 1024
num_hiddens2 = 1024
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 预处理文件所需参数
file_batch_size = 192
# 总共约八万条数据集
# 数据集采用 70:7:3 的比例划分训练集、验证集和测试集
# 数据集的长度
train_len, test_len, val_len = 0, 0, 0
# 通过random列表确定概率
random_list = [0,1,2,3,4,5,6,7,8,9]
# 文件根目录
root_path = '.'

数据预处理

因为自己家电脑配置比较渣,RTX3060 6G显存 + 16G DDR4内存,如果采用网上大部分类似项目采用的数据预处理方法,一是电脑带不动,显卡压力太大,二是内存压力太大,导致训练的时候内存和显卡的性能严重的影响了训练的速度。

我采用的方法是预处理文件先将图片处理为tensor矩阵,数据集约八万张图片,处理后约40G的tensor,然后再在训练的时候直接读取tensor而不需要处理文件。这么做会带来另外一个问题,存取为.pt文件时,先将文件存在内存中,作为列表的形式存在,然后再存取;读取.pt文件中的矩阵时,将矩阵读取到内存中,再用迭代器进行迭代并送到train函数中进行运算。这样的话所有的tensor矩阵都会进入内存中,导致我这16G内存(实际上只有8G windows系统都占用了一半)远远不够储存40G的训练集。

于是我决定采取分块存取,分块取用的方法,主要是训练集太多导致内存吃不住,但是验证集和测试集内存是几乎毫无压力的,所以将训练集分为n个模块(实际训练测试后分为了18个模块),经测试,内存占用率未达到上限,内存不会影响训练的速度,训练时,显卡的 CUDA 加速占用率几乎时刻为100%。预处理数据时也同理,和读出数据时的内存情况差不多。

训练时的内存情况

随机化

分别在几个地方进行了随机化

  • 存取数据时,利用 random.shuffle(list) 函数随机化
  • 读取数据时,利用Dataloader中的 shuffle = True 随机化迭代器
  • 预处理图像时经过裁剪, 放缩, 加高斯噪声的方式随机化
  • 迭代器分块读取时随机化读取每个文件 (文件0-17共18个文件随机抽取一个作为迭代对象)
        #预处理时的随机化
        # 训练集的图片处理
        self.trian_transform = transform.Compose(
            [
            # 缩放图片为224*224
            transform.RandomResizedCrop(224),
            # 数据增强
            # 随机水平翻转
            transform.RandomHorizontalFlip(),
            # 随机色彩
            transform.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.3),
            # 随机竖直翻转
            # transform.RandomVerticalFlip(),
            # 转化为 channel * height * width 
            transform.ToTensor(),
            transform.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ]
        )
        # 测试集和验证集的图片处理
        self.val_transform = transform.Compose(
            [
            transform.Resize(224),
            transform.ToTensor()
            ]
        )

数据读取/进行训练

采用Adam优化

因为采用了分块读取,因此,采用分块更新迭代器的方式进行数据存取

# main.py
    # 初始化数据
    # 更新train_iter
    def update_train_iter(train_mod_num):
        logger.info('update train_iter ' + str(train_mod_num))
        train_dataset = waste_sorting_dataset(train_mod_num)
        train_loader = DataLoader(dataset = train_dataset, batch_size = batch_size, shuffle = True, pin_memory = True, num_workers = num_workers)                     
        return train_loader
    # 初始化验证集的迭代器
    def init_val_iter():
        val_data = waste_sorting_dataset(-2)
        val_iter = DataLoader(dataset = val_data, batch_size = batch_size, pin_memory = True, num_workers = num_workers)
        logger.info('nums of val_data: '+str(len(val_data)))
        return val_iter
    # 初始化测试集的迭代器
    def init_test_iter():
        test_data = waste_sorting_dataset(-1)
        test_iter = DataLoader(dataset = test_data, batch_size = batch_size, pin_memory = True, num_workers = num_workers)
        logger.info('nums of test_iter: '+ str(len(test_data)))
        return test_iter


# waste_sorting_dataset.py
from torch.utils.data import Dataset, DataLoader
import random
import torch

#-2为验证模式, -1为测试模式, 0以上为训练模式
class waste_sorting_dataset(Dataset):
    def __init__(self, train_mod_num):
        if train_mod_num == -2:
            self.imgs = torch.load('./pt/val.pt')
        elif train_mod_num == -1:
            self.imgs = torch.load('./pt/test.pt')
        else:
            self.imgs = torch.load('./pt/train_' + str(train_mod_num) + '.pt')
    def __getitem__(self, index):
        img, label = self.imgs[index]
        return img, label
    def __len__(self):
        return len(self.imgs)

有别于传统的类似做法,采用了传递函数而非传递迭代器实例化对象的方式控制内存的占用,训练函数就非常正常,按照传统做法写的train函数,只是增加了一层外层循环便于迭代器的循环更新,减小内存的占用。

# train.py
from torch import gather
from loguru import logger
import torch
import time
import random
from draw import semilogy

# 开始训练
# 传入函数而非对象 减少内存的占用 用的时候再实例化
def train(net, update_train_iter, init_val_iter, optimizer, scheduler, device, num_epochs)

最后将训练数据保存至log文件中便于查看以及画图分析


海纳百川 有容乃大