diff --git a/README.zh.md b/README.zh.md index f16a5bd6607e62a31b3b9850c465547a44269dbd..04af0c2cc31f3f1889145633cd6e3e752ba03a69 100644 --- a/README.zh.md +++ b/README.zh.md @@ -56,46 +56,90 @@ Apex及Apex-patch的卸载只需执行命令: ## 自动混合精度(AMP) -使用apex.amp进行混合精度训练: -``` -model = torch.nn.Linear(D_in, D_out).npu() -optimzier = torch.optim.SGD(model.parameters(), lr=1e-3) - -model, optimizer = amp.initialize(model, optimizer, opt_level='O1') -... -with amp.scale_loss(loss, optimizer) as scaled_loss: -  scaled_loss.backward() -... +导入apex包,使用apex.amp.initialize方法包装模型和优化器并指定opt_level为'O1',并使用apex.amp.scale_loss方法包装损失值和优化器对损失进行动态缩放,实现混合精度训练: +```diff + import torch ++ import apex # 导入apex包 + + + epoch = 10 # 训练轮次 + loss_function = torch.nn.MSELoss() # 损失函数 + input_data = [torch.rand(16, device='npu') for i in range(10)] # 构造简易模型输入 + dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data] # 构造简易数据集 + model = torch.nn.Linear(in_features=16, out_features=1, device='npu') # 构造简易模型 + optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3) # 创建优化器 ++ # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练,opt_level取'O0'时,表示使用FP32训练,opt_level取'O1'或'O2'时,表示使用混合精度训练 ++ model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1') + for i in range(epoch): + for x, target in dataset: + optimizer.zero_grad() + output = model(x) + loss = loss_function(output, target) +- loss.backward() ++ # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放 ++ with apex.amp.scale_loss(loss, optimizer) as scaled_loss: ++ scaled_loss.backward() + optimizer.step() ``` 详细使用方式请请参考[Apex官方文档](https://nvidia.github.io/apex/amp.html)。 ## 使用融合梯度进行scale/unscale优化加速 -在amp.initialize()中将参数combine_grad设置为True,如: -``` -model = torch.nn.Linear(D_in, D_out).npu() -optimzier = torch.optim.SGD(model.parameters(), lr=1e-3) - -model, optimizer = amp.initialize(model, optimizer, opt_level='O1', combine_grad=True)  # 增加combine_grad参数 -... -with amp.scale_loss(loss, optimizer) as scaled_loss: -  scaled_loss.backward() -... +导入apex包,使用apex.amp.initialize方法包装模型和优化器并将参数combine_grad设置为True,示例如下: +```diff + import torch ++ import apex # 导入apex包 + + + epoch = 10 # 训练轮次 + loss_function = torch.nn.MSELoss() # 损失函数 + input_data = [torch.rand(16, device='npu') for i in range(10)] # 构造简易模型输入 + dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data] # 构造简易数据集 + model = torch.nn.Linear(in_features=16, out_features=1, device='npu') # 构造简易模型 + optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3) # 创建优化器 ++ # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练,opt_level='O1'表示使用混合精度训练,combine_grad=True表示开启融合梯度 ++ model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1', combine_grad=True) + for i in range(epoch): + for x, target in dataset: + optimizer.zero_grad() + output = model(x) + loss = loss_function(output, target) +- loss.backward() ++ # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放 ++ with apex.amp.scale_loss(loss, optimizer) as scaled_loss: ++ scaled_loss.backward() + optimizer.step() ``` ## 使用融合优化器优化加速 -将torch原生优化器torch.optim.xxx替换为apex.optimizers.xxx, 其中xxx为融合优化器名称,apex-patch支持的优化器见*特性介绍*。 -``` -model = torch.nn.Linear(D_in, D_out).npu() -optimzier = apex.optimizers.NpuFusedSGD(model.parameters(), lr=1e-3) # 使用apex.optimizers.NpuFusedSGD - -model, optimizer = amp.initialize(model, optimizer, opt_level='O1', combine_grad=True) -... -with amp.scale_loss(loss, optimizer) as scaled_loss: -  scaled_loss.backward() -... +导入apex包,使用apex.amp.initialize方法包装模型和优化器并将参数combine_grad设置为True,并将torch原生优化器torch.optim.xxx替换为apex.optimizers.xxx, 其中xxx为融合优化器名称,示例如下所示,apex-patch支持的优化器见*特性介绍*。 +```diff + import torch ++ import apex # 导入apex包 + + + epoch = 10 # 训练轮次 + loss_function = torch.nn.MSELoss() # 损失函数 + input_data = [torch.rand(16, device='npu') for i in range(10)] # 构造简易模型输入 + dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data] # 构造简易数据集 + model = torch.nn.Linear(in_features=16, out_features=1, device='npu') # 构造简易模型 +- optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3) # 创建优化器 ++ # 使用apex.optimizers.NpuFusedSGD替换torch.optim.SGD ++ optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=1e-3) ++ # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练,opt_level='O1'表示使用混合精度训练,combine_grad=True表示开启融合梯度 ++ model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1', combine_grad=True) + for i in range(epoch): + for x, target in dataset: + optimizer.zero_grad() + output = model(x) + loss = loss_function(output, target) +- loss.backward() ++ # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放 ++ with apex.amp.scale_loss(loss, optimizer) as scaled_loss: ++ scaled_loss.backward() + optimizer.step() ``` # 特性介绍