TopK算子返回的全零的Tensor的解决

经验分享 Tech Blogs 模型训练-Model Training
1

1 系统环境

硬件环境(Ascend/GPU/CPU): GPU
MindSpore版本: mindspore=1.10.1
执行模式（PyNative/ Graph）:Graph
Python版本: Python=3.9.16
操作系统平台: Linux Ubuntu 20.04

2 报错信息

2.1 问题描述

在使用ops.top_k算子进行计算时（详见下方用例），会出现计算结果错误。
预期结果
如果正确执行，输出应该类似于如下结果：

$ python3 run_topk.py  
26756342

实际执行结果如下，返回的是一个全零的Tensor：

$ python3 run_topk.py  
0

备注：经过一定数值范围的测试，发现在K值选取1~128范围之间的返回结果有可能是正确的，或者两个打印模块同时打印输出结果也是正确的，唯独单次打印K值大于128的情况，返回结果异常。

2.2 报错信息

计算结果不符合预期。

2.3 脚本代码（代码格式，可上传附件）

# run_topk.py  
import os  
os.environ['GLOG_v']='4'  
import numpy as np  
import mindspore as ms  
from mindspore import Tensor, ops, nn  
    
    
class TopKNet(nn.Cell):  
    def __init__(self):  
        super().__init__()  
        self.top_k = ops.top_k  
    
    def construct(self, x, num):  
        x, ids = self.top_k(x, num)  
        return ids  
    
np.random.seed(1)  
input_x = Tensor(np.random.random((1, 648, 648)), ms.float16)  
print (TopKNet()(input_x, 128).sum())  
# print (TopKNet()(input_x, 129).sum())

直接运行上述用例，得到一个大的整数结果，这是正确场景；
注释第一个打印输出，取消第二个打印输出的注释，也就是选取前129个最大的值；
再次执行案例，得到的结果是0，或者是一个异常值。

3 根因分析

这是因为在mindspore关于topk的cuda代码中，一个核函数在做数据交换的时候，发生了数据类型强转(float和int)。导致当index的数值较大时(7 * 2388 * 2388+289)，会出现精度损失。

4 解决方案

第一个办法，可以将mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/topk_lib.cuh这个位置下的58行进行修改

修改前

T b_tmp = b[i_1];

修改后

S b_tmp = b[i_1];

第二个办法是使用2.0版本的，这个更方便可靠。