windows 安装apex_Nvidia Apex安装

降低训练精度提高batch大小,512*512的图片在8G的显存上batch只能为1,2,使用了apex后可以增大到10。只支持pytorch。

官网地址:NVIDIA/apex

官方说支持Ubuntu系统,Windows只是实验性的,但是我在win10上也安装成功了。

前提

已安装好CUDA和CUDNN,且版本适配。

2.在安装前先检查一下,电脑的cuda版本和pytorch内的cuda版本是否一样,不一样的话就把低版本的进行升级。

# 查看电脑的cuda版本、

>> nvcc -V

# pytorch内的cuda版本

import torch

torch.version.cuda

安装

按照官网的命令输入即可

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

$ cd apex

$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

如果报错出现什么pip 137..我不记得了

在保证cuda版本一致的前提下进入apex文件夹,使用命令:python setup.py install 即可安装成功

验证安装成功

from apex import amp

没有报错就是成功了。

Apex的使用

net = xxxNet()

net.cuda()

net.train()

params_low_lr = []

params_high_lr = []

for n, p in net.named_parameters():

if 'encoder' in n:

params_low_lr.append(p)

else:

params_high_lr.append(p)

opt = Adam([{'params': params_low_lr, 'lr': 5e-5},

{'params': params_high_lr, 'lr': 1e-4}], weight_decay=settings.WEIGHT_DECAY)

net, opt = amp.initialize(net, opt, opt_level="O1") # 这里要添加这句代码

...

...

loss = ...

if not torch.isnan(loss):

opt.zero_grad()

with amp.scale_loss(loss, opt) as scaled_loss:

scaled_loss.backward() # loss要这么用

opt.step()

当中断后继续训练时

# Initialization

opt_level = 'O1'

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

# Train your model

...

# Save checkpoint

checkpoint = {

'model': model.state_dict(),

'optimizer': optimizer.state_dict(),

'amp': amp.state_dict()

}

torch.save(checkpoint, 'amp_checkpoint.pt')

...

# Restore

model = ...

optimizer = ...

checkpoint = torch.load('amp_checkpoint.pt')

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

model.load_state_dict(checkpoint['model']) # 注意,load模型需要在amp.initialize之后!!!

optimizer.load_state_dict(checkpoint['optimizer'])

amp.load_state_dict(checkpoint['amp'])

# Continue training

···

用户指定数据格式

amp.initialize(net, opt, opt_level=“O1”)

其中的opt-level参数是用户指定采用何种数据格式做训练的。

O0:纯FP32训练,可以作为accuracy的baseline;

O1:混合精度训练(推荐使用),根据黑白名单自动决定使用FP16(GEMM, 卷积)还是FP32(Softmax)进行计算。

O2:“几乎FP16”混合精度训练,不存在黑白名单,除了Batch norm,几乎都是用FP16计算。

O3:纯FP16训练,很不稳定,但是可以作为speed的baseline;

溢出问题

相乘的操作,16位*16位=32位,精度不会有大的影响,但16位+32位由于后面32位的数太小,会进行舍入只得到一个16位数,影响较大,所以当涉及到“加法”时最好恢复到32位精度,比如 sigmoid,softmax。

我们仅需在模型定义中,在构造函数init中的某一个位置。加上下面这段:

from apex import amp

class xxxNet(Module):

def __init__(using_map=False)

...

...

if using_amp:

amp.register_float_function(torch, 'sigmoid')

amp.register_float_function(torch, 'softmax')

用register_float_function指明后面的函数需要使用float类型。注意第二实参是string类型

和register_float_function相似的注册函数还有

amp.register_half_function(module, function_name)

amp.register_float_function(module, function_name)

amp.register_promote_function(module, function_name)

你必须在使用amp.initialize之前使用注册函数,所以最好的位置就放在模型的构造函数中

你可能感兴趣的:(windows,安装apex)