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