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

训练准备概览

    # ./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()
            ]
        )

数据读取/进行训练

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

# 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文件中便于查看以及画图分析