单精度和半精度混合训练
單精度和半精度混合訓練
概述
混合精度訓練方法,通過混合使用單精度和半精度數據格式,加速深度神經網絡訓練的過程,同時保持了單精度訓練所能達到的網絡精度。混合精度訓練能夠加速計算過程,同時減少內存使用和存取,并使得在特定的硬件上可以訓練更大的模型或batch size。
對于FP16的算子,若給定的數據類型是FP32,MindSpore框架的后端會進行降精度處理。用戶可以開啟INFO日志,并通過搜索關鍵字“Reduce precision”查看降精度處理的算子。
計算流程
MindSpore混合精度典型的計算流程如下圖所示:
- 參數以FP32存儲;
- 正向計算過程中,遇到FP16算子,需要把算子輸入和參數從FP32 cast成FP16進行計算;
- 將Loss層設置為FP32進行計算;
- 反向計算過程中,首先乘以Loss Scale值,避免反向梯度過小而產生下溢;
- FP16參數參與梯度計算,其結果將被cast回FP32;
- 除以Loss scale值,還原被放大的梯度;
- 判斷梯度是否存在溢出,如果溢出則跳過更新,否則優化器以FP32對原始參數進行更新。
本文通過自動混合精度和手動混合精度的樣例來講解計算流程。
自動混合精度
使用自動混合精度,需要調用相應的接口,將待訓練網絡和優化器作為輸入傳進去;該接口會將整張網絡的算子轉換成FP16算子(除BatchNorm算子和Loss涉及到的算子外)。可以使用amp接口和Model接口兩種方式實現混合精度。
使用amp接口具體的實現步驟為: - 引入MindSpore的混合精度的接口amp;
- 定義網絡:該步驟和普通的網絡定義沒有區別(無需手動配置某個算子的精度);
- 使用amp.build_train_network接口封裝網絡模型、優化器和損失函數,設置level參數,參考https://www.mindspore.cn/doc/api_python/zh-CN/r1.1/mindspore/mindspore.html#mindspore.build_train_network。在該步驟中,MindSpore會將有需要的算子自動進行類型轉換。
代碼樣例如下:
import numpy as np
import mindspore.nn as nn
from mindspore import Tensor, context
import mindspore.ops as ops
from mindspore.nn import Momentum
The interface of Auto_mixed precision
from mindspore import amp
context.set_context(mode=context.GRAPH_MODE)
context.set_context(device_target=“Ascend”)
Define network
class Net(nn.Cell):
def init(self, input_channel, out_channel):
super(Net, self).init()
self.dense = nn.Dense(input_channel, out_channel)
self.relu = ops.ReLU()
def construct(self, x):x = self.dense(x)x = self.relu(x)return x
Initialize network
net = Net(512, 128)
Define training data, label
predict = Tensor(np.ones([64, 512]).astype(np.float32) * 0.01)
label = Tensor(np.zeros([64, 128]).astype(np.float32))
Define Loss and Optimizer
loss = nn.SoftmaxCrossEntropyWithLogits()
optimizer = Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
train_network = amp.build_train_network(net, optimizer, loss, level=“O3”, loss_scale_manager=None)
Run training
output = train_network(predict, label)
使用Model接口具體的實現步驟為:
- 引入MindSpore的模型訓練接口Model;
- 定義網絡:該步驟和普通的網絡定義沒有區別(無需手動配置某個算子的精度);
- 創建數據集。該步驟可參考 https://www.mindspore.cn/tutorial/training/zh-CN/r1.1/use/data_preparation.html;
- 使用Model接口封裝網絡模型、優化器和損失函數,設置amp_level參數,參考https://www.mindspore.cn/doc/api_python/zh-CN/r1.1/mindspore/mindspore.html#mindspore.Model。在該步驟中,MindSpore會將有需要的算子自動進行類型轉換。
代碼樣例如下:
import numpy as np
import mindspore.nn as nn
from mindspore.nn.metrics import Accuracy
from mindspore import context, Model
from mindspore.common.initializer import Normal
from src.dataset import create_dataset
context.set_context(mode=context.GRAPH_MODE)
context.set_context(device_target=“Ascend”)
Define network
class LeNet5(nn.Cell):
“”"
Lenet network
Args:num_class (int): Number of classes. Default: 10.num_channel (int): Number of channels. Default: 1.Returns:Tensor, output tensor
Examples:>>> LeNet(num_class=10)"""
def __init__(self, num_class=10, num_channel=1):super(LeNet5, self).__init__()self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')self.fc1 = nn.Dense(16 * 5 * 5, 120, weight_init=Normal(0.02))self.fc2 = nn.Dense(120, 84, weight_init=Normal(0.02))self.fc3 = nn.Dense(84, num_class, weight_init=Normal(0.02))self.relu = nn.ReLU()self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)self.flatten = nn.Flatten()def construct(self, x):x = self.max_pool2d(self.relu(self.conv1(x)))x = self.max_pool2d(self.relu(self.conv2(x)))x = self.flatten(x)x = self.relu(self.fc1(x))x = self.relu(self.fc2(x))x = self.fc3(x)return x
create dataset
ds_train = create_dataset("/dataset/MNIST/train", 32)
Initialize network
network = LeNet5(10)
Define Loss and Optimizer
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction=“mean”)
net_opt = nn.Momentum(network.trainable_params(),learning_rate=0.01, momentum=0.9)
model = Model(network, net_loss, net_opt, metrics={“Accuracy”: Accuracy()}, amp_level=“O3”)
Run training
model.train(epoch=10, train_dataset=ds_train)
手動混合精度
MindSpore還支持手動混合精度。假定在網絡中只有一個Dense Layer要用FP32計算,其他Layer都用FP16計算。混合精度配置以Cell為粒度,Cell默認是FP32類型。
以下是一個手動混合精度的實現步驟:
- 定義網絡:該步驟與自動混合精度中的步驟2類似;
- 配置混合精度:通過net.to_float(mstype.float16),把該Cell及其子Cell中所有的算子都配置成FP16;然后,將模型中的dense算子手動配置成FP32;
- 使用TrainOneStepCell封裝網絡模型和優化器。
代碼樣例如下:
import numpy as np
import mindspore.nn as nn
from mindspore import dtype as mstype
from mindspore import Tensor, context
import mindspore.ops as ops
from mindspore.nn import WithLossCell, TrainOneStepCell
from mindspore.nn import Momentum
context.set_context(mode=context.GRAPH_MODE)
context.set_context(device_target=“Ascend”)
Define network
class Net(nn.Cell):
def init(self, input_channel, out_channel):
super(Net, self).init()
self.dense = nn.Dense(input_channel, out_channel)
self.relu = ops.ReLU()
def construct(self, x):x = self.dense(x)x = self.relu(x)return x
Initialize network
net = Net(512, 128)
Set mixing precision
net.to_float(mstype.float16)
net.dense.to_float(mstype.float32)
Define training data, label
predict = Tensor(np.ones([64, 512]).astype(np.float32) * 0.01)
label = Tensor(np.zeros([64, 128]).astype(np.float32))
Define Loss and Optimizer
loss = nn.SoftmaxCrossEntropyWithLogits()
optimizer = Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepCell(net_with_loss, optimizer)
train_network.set_train()
Run training
output = train_network(predict, label)
約束
使用混合精度時,只能由自動微分功能生成反向網絡,不能由用戶自定義生成反向網絡,否則可能會導致MindSpore產生數據格式不匹配的異常信息。
總結
以上是生活随笔為你收集整理的单精度和半精度混合训练的全部內容,希望文章能夠幫你解決所遇到的問題。