『PyTorch』屌丝的PyTorch玩法

Posted 2021-02-21 ice-coder

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

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

2. 安装apex

cd apex
python setup.py install

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

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

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

4. 测试安装成功

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，注意，是欧不是零