dsmm_test.py 文件源码

def test_backward_batch():
    i = torch.LongTensor([[0, 0, 0, 1, 1, 1],
                          [0, 1, 1, 0, 1, 1],
                          [2, 0, 2, 2, 0, 2]])
    v = torch.FloatTensor([3, 4, 5, 6, 7, 8])
    sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 3]))
    dense = Variable(torch.randn(2, 3, 4), requires_grad=True)
    dense_copy = Variable(dense.data.clone(), requires_grad=True)
    grad_output = torch.randn(2, 2, 4)

    res = gpytorch.dsmm(Variable(sparse), dense)
    res.backward(grad_output)
    actual = torch.matmul(Variable(sparse.to_dense()), dense_copy)
    actual.backward(grad_output)
    assert(torch.norm(dense.grad.data - dense_copy.grad.data) < 1e-5)