粗暴啤酒

文章
4
资源
0
加入时间
3年0月28天

Pytorch 源码浅析1.Topk(上)简介

简介本章节将介绍笔者学习pytorch源码中遇到的一些并行算法(基于pytorch)。当然这些源码都可以在官方的github上获取。入口整体思路入口为gather_topk函数,首先会经过RadixSelect函数找到第K大小的数(可能不是唯一的)。然后再并行找到大于这个值的所有数并按序前插以确定Index(本节举例最大Topk)。首先需要了解的前提知识为基数排序,在Pytorch中topk是基于基数排序的变体实现的,了解其原理有助于了解topk的实现。基数排序RadixSelect说这个函