『PyTorch』屌丝的PyTorch玩法

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 安装

克隆源代码

git clone https://github.com/NVIDIA/apex

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

安装apex

cd apex
python setup.py install

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

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

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

测试安装成功

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

本站文章如无特殊说明，均为本站原创，如若转载，请注明出处：『PyTorch』屌丝的PyTorch玩法 - Python技术站

『PyTorch』屌丝的PyTorch玩法

1. prefetch_generator

2. Apex

2.1 安装

2.2 使用

相关文章