modules.py 文件源码

python
阅读 44 收藏 0 点赞 0 评论 0

项目:pytorch_workplace 作者: DingKe 项目源码 文件源码
def topk(x, k):
    ''' Proof-of-concept implementation of topk.
    '''
    original_size = None
    if x.dim() > 2:
        original_size = x.size()
        x = x.view(x.size(0), -1)
    ax = torch.abs(x.data)
    topk, _ = ax.topk(k)
    topk = topk[:, -1]
    y = x.clone()
    # zero out small values
    y[ax < topk.repeat(x.size(-1), 1).transpose(0, 1)] = 0

    if original_size:
        y = y.view(original_size)
    return y
评论列表
文章目录


问题


面经


文章

微信
公众号

扫码关注公众号