1. prefetch_generator

使用 prefetch_generator库 在后台加载下一batch的数据,原本PyTorch默认的DataLoader会创建一些worker线程来预读取新的数据,但是除非这些线程的数据全部都被清空,这些线程才会读下一批数据。使用prefetch_generator,我们可以保证线程不会等待,每个线程都总有至少一个数据在加载。

  • 安装

    pip install prefetch_generator
    
  • 使用
    之前加载数据集的正确方式是使用torch.utils.data.DataLoader,现在我们只要利用这个库,新建个DataLoaderX类继承DataLoader并重写__iter__方法即可

    from torch.utils.data import DataLoader
    from prefetch_generator import BackgroundGenerator
    
    class DataLoaderX(DataLoader):
    
        def __iter__(self):
            return BackgroundGenerator(super().__iter__())
        
    

    之后这样用:

    train_dataset = MyDataset(".........")
    train_loader = DataLoaderX(dataset=train_dataset, 
                               batch_size=batch_size, num_workers=4, shuffle=shuffle)
        
    

2. Apex

2.1 安装

  1. 克隆源代码
git clone https://github.com/NVIDIA/apex

可以先下载到码云,再下载到本地

  1. 安装apex
cd apex
python setup.py install

最好打开PyCharm的终端进行安装,这样实在Anaconda的环境里安装了

  1. 删除刚刚clone下来的apex文件夹,然后重启PyCharm

【注意】安装PyTorch和cuda时注意版本对应,要按照正确流程安装

  1. 测试安装成功
from apex import amp

如果导入不报错说明安装成功

2.2 使用

from apex import amp  # 这个必须的,其他的导包省略了

train_dataset = MyDataset("......")
train_loader = DataLoader(dataset=train_dataset, batch_size=2, num_workers=4, shuffle=True)

model = MyNet().to(device)  # 创建模型

criterion = nn.MSELoss()  # 定义损失函数

optimizer = optim.Adam(net.parameters(), lr=learning_rate, weight_decay=0.00001)  # 优化器

net, optimizer = amp.initialize(net, optimizer, opt_level="O1")  # 这一步很重要

# 学习率衰减
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer=optimizer, mode="min",factor=0.1, patience=3, 
    verbose=False,cooldown=0, min_lr=0.0, eps=1e-7)

for epoch in range(epochs):
    net.train()  # 训练模式

    train_loss_epoch = []  # 记录一个epoch内的训练集每个batch的loss
    test_loss_epoch = []  # 记录一个epoch内测试集的每个batch的loss

    for i, data in enumerate(train_loader):
        # forward
        x, y = data
        x = x.to(device)
        y = y.to(device)

        outputs = net(x)

        # backward
        optimizer.zero_grad()
        
        loss = criterion(outputs, labels)  
        
        # 这一步也很重要
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()

        # 更新权重
        optimizer.step()

    scheduler.step(1)  # 更新学习率。每1步更新一次

  • 主要是添加了三行代码
  • scaled_loss 是将原loss放大了,所以要保存loss应该保存之前的值,这种放大防止梯度消失

考察amp.initialize(net, optimizer, opt_level="O1")opt_level参数

  • opt_level=O0(base)
    表示的是当前执行FP32训练,即正常的训练

  • opt_level=O1(推荐)
    表示的是当前使用部分FP16混合训练

  • opt_level=O2

    表示的是除了BN层的权重外,其他层的权重都使用FP16执行训练

  • opt_level=O3
    表示的是默认所有的层都使用FP16执行计算,当keep_batch norm_fp32=True,则会使用cudnn执行BN层的计算,该优化等级能够获得最快的速度,但是精度可能会有一些较大的损失

一般我们用O1级别就行,最多O2,注意,是不是