MindSpore在昇腾Ascend平台使用全精度计算
背景信息
在昇腾Ascend芯片上，有不少算子为了能够拥有更高运行速度而使得数据类型只支持FP16；也有不少算子同时支持FP16与FP32数据类型，但对于某些网络来说，FP16数据类型满足不了精度要求；针对在Ascend环境上只支持FP16数据类型的算子须通过异构计算的方法将算子异构到CPU运行，从而实现由FP16到FP32数据类型的转换；针对在Ascend环境上同时支持FP16与FP32数据类型的算子则可通过to_float()方法按照用户需求进行转换。
1、异构计算
1.1、示例代码
实现功能：使用异构方式将特定卷积层由FP16转化为精度较高的FP32类型
import mindspore.nn as nnimport numpy as npfrom mindspore import Tensorfrom mindspore import contextimport mindspore.ops as opsimport mindspore
#设置程序运行环境，并保存程序运行IR图
context.set_context(device_target="Ascend",save_graphs=True, save_graphs_path="/home/ma-user/xxx/ir_save/")
class PrintDemo(nn.Cell):
def __init__(self):
super(PrintDemo, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=4, stride=1 ,has_bias=False, weight_init='normal', pad_mode='valid')
self.bn1 = nn.BatchNorm2d(6)
self.relu = nn.ReLU()
self.conv2 = nn.Conv2d(in_channels=6, out_channels=2, kernel_size=2, pad_mode="valid")
self.conv2.conv2d.add_prim_attr('primitive_target', 'CPU')
self.conv3 = nn.Conv2d(in_channels=2, out_channels=6, kernel_size=2, pad_mode="valid")
def construct(self, input_x):
x = self.conv1(input_x)
x = self.bn1(x)
x = self.relu(x)
x = self.conv2(x)
x = self.conv3(x)
return x
def test():
input_x = Tensor(np.ones([1, 1, 32, 32]), mindspore.float32)
net = PrintDemo()
result = net(input_x)
return result
total_result = test()
print("total_result:",total_result)异构计算关键代码解析：
步骤一：用户对训练网络硬件环境进行设置
from mindspore import context
context.set_context(device_target="Ascend")步骤二：用户设置特定算子到特定硬件环境运行
prim.add_prim_attr("primitive_target", "CPU")1.2、运行结果分析
程序运行结束后，在IR图中查看特定算子是否异构成功，以示例代码为例进行分析：
数据的流向为：Ascend（conv1、bn1、relu）—>CPU（conv2）—>Ascend（conv3），MindSpore在不同的硬件环境会编译生成不同的执行图，所以该示例会生成3张执行图，即对应IR图中会生成3份trace_xxx文件。
打开在Ascend执行的第一张图所对应的trace_code_graph_xxx文件，此时数据流经过第一个卷积层所对应的数据类型为：float32(输入)—>float16(计算)—>float32(输出)
打开在CPU执行的第二张图所对应的trace_code_graph_xxx文件，此时数据流经过第二个卷积层所对应的数据类型为：float32(输入)—>float32(计算)—>float32(输出)
打开在Ascend执行的第三张图所对应的trace_code_graph_xxx文件，此时数据流经过第三个卷积层所对应的数据类型为：float32(输入)—>float16(计算)—>float32(输出)
注：IR图可通过对context.set_context函数中save_graphs=True以及save_graphs_path="/path/进行获取IR 图。
2、to_float方法
2.1、示例代码
import mindspore.nn as nnimport numpy as npfrom mindspore import Tensorfrom mindspore import contextimport mindspore.ops as opsimport mindspore from mindspore import dtype as mstype
context.set_context(device_target="Ascend",save_graphs=True, save_graphs_path="/home/ma-user/xxx/ir_save1/")
class PrintDemo(nn.Cell):
def __init__(self):
super(PrintDemo, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=4, stride=1 ,has_bias=False, weight_init='normal', pad_mode='valid')
self.bn1 = nn.BatchNorm2d(6)
self.relu = nn.ReLU()
self.dense = nn.Dense(in_channels=29,out_channels=6)
self.conv2 = nn.Conv2d(in_channels=6, out_channels=2, kernel_size=4, stride=1 ,has_bias=False, weight_init='normal', pad_mode='valid')
def construct(self, input_x):
x = self.conv1(input_x)
x = self.bn1(x)
x = self.relu(x)
x = self.dense(x)
x = self.conv2(x)
return x
def test():
input_x = Tensor(np.ones([1, 1, 32, 32]), mindspore.float32)
net = PrintDemo()
net.to_float(mindspore.float16)
net.dense.to_float(mindspore.float32)
result = net(input_x)
return result
total_result = test()
print("total_result:",total_result.shape)to_float方法关键代码解析：
net.to_float(mindspore.float16)#将网络中所有数据转换为float16数据类型
net.dense.to_float(mindspore.float16)#将网络中Dense层中的计算所用数据的数据类型转换为float16
用户可根据需求对特定算子转换到所需要的数据类型；
2.2、运行结果分析
程序运行结束后，可通过IR图查看特定算子是否异构成功，以示例代码为例进行分析。
打开所保存的IR图路径下trace_code_graph_xxx文件，此时数据流经过第一个卷积层conv1所对应的数据类型为：float16(输入)—>float16（计算）—>float16（输出）
在trace_code_graph_xxx文件中的全连接层nn.Dense使用to_float方法后，此时数据流经过在该全连接层的所有计算全已都转换为FP32数据类型；此时数据流的数据类型为：float32(输入)—>float32(计算)—>float32(输出)
在最后一层的卷积层conv2中，此时数据流经过第二个卷积层conv2所对应的数据类型为：float16(输入)—>float16(计算)—>float16(输出)
