给定一维张量中的舍入值，舍入二维张量的正确实现是什么？

Question

这是我到目前为止所做的：

def round_values(predictions):
    # Rounding values
    rounded = torch.tensor([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0])

    # Find nearest rounding value 
    dif = predictions.unsqueeze(2) - rounded.view(1,1,-1)
    indexes = torch.argmin(abs(dif), dim=2)

    # Fill tensor with nearest round value
    # TO IMPROVE
    rounded_predictions = torch.zeros(*predictions.shape)
    for i in range(rounded_predictions.shape[0]):
        for j in range(rounded_predictions.shape[1]):
            index = indexes[i,j]
            rounded_predictions[i,j] = rounded[index]
        
    return rounded_predictions 

a = 4*torch.rand((6,6)) + 1
print(a)
print(round_values(a))

双环太可怕了。 我想要类似rounded [indexes]的东西，所以它返回一个与具有rounded值的索引具有相同形状的张量，但我还没有找到一种方法来在张量级别而不是在元素方面做到这一点。 我想知道吗？

Answer 1

100 x 100 阵列（在 cpu 上）快一千倍

def round_values_bob(predictions):
    return torch.clamp(torch.round(predictions * 2) / 2, 1, 5)