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
评论列表
文章目录