modules.py 文件源码

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