carpet/proximity_tv.py
def backward(ctx, grad_output):
        """Compute the gradient of proxTV using implicit gradient."""
        batch_size, n_dim = grad_output.shape
        sign_z, = ctx.saved_tensors
        device = grad_output.device
        S = sign_z != 0
        S[:, 0] = True
        sign_z[:, 0] = 0
        # XXX do clever computations
        L = torch.triu(torch.ones((n_dim, n_dim), dtype=torch.float64,
                       device=device))

        grad_x, grad_lbda = [], []
        for i in range(batch_size):
            L_S = L[:, S[i]]                     # n_dim x |S|
            grad_u = grad_output[i].matmul(L_S)  # 1 x |S|
            H_S = torch.inverse(L_S.t().matmul(L_S))
            grad_x.append(grad_u.matmul(H_S.matmul(L_S.t())))
            grad_lbda.append(grad_u.matmul(H_S.matmul(-sign_z[i][S[i]])))
        grad_x = torch.stack(grad_x)
        grad_lbda = torch.stack(grad_lbda)
        return (grad_x, grad_lbda)
Compute the gradient of proxTV using implicit gradient.
Compute the gradient of proxTV using implicit gradient.
Compute the gradient of proxTV using implicit gradient .
def backward(ctx, grad_output):
        batch_size, n_dim = grad_output.shape
        sign_z, = ctx.saved_tensors
        device = grad_output.device
        S = sign_z != 0
        S[:, 0] = True
        sign_z[:, 0] = 0
        L = torch.triu(torch.ones((n_dim, n_dim), dtype=torch.float64,
                       device=device))
        grad_x, grad_lbda = [], []
        for i in range(batch_size):
            L_S = L[:, S[i]]                     
            grad_u = grad_output[i].matmul(L_S)  
            H_S = torch.inverse(L_S.t().matmul(L_S))
            grad_x.append(grad_u.matmul(H_S.matmul(L_S.t())))
            grad_lbda.append(grad_u.matmul(H_S.matmul(-sign_z[i][S[i]])))
        grad_x = torch.stack(grad_x)
        grad_lbda = torch.stack(grad_lbda)
        return (grad_x, grad_lbda)
def backward ( ctx , grad_output ) : batch_size , n_dim = grad_output . shape sign_z , = ctx . saved_tensors device = grad_output . device S = sign_z != 0 S [ : , 0 ] = True sign_z [ : , 0 ] = 0 L = torch . triu ( torch . ones ( ( n_dim , n_dim ) , dtype = torch . float64 , device = device ) ) grad_x , grad_lbda = [ ] , [ ] for i in range ( batch_size ) : L_S = L [ : , S [ i ] ] grad_u = grad_output [ i ] . matmul ( L_S ) H_S = torch . inverse ( L_S . t ( ) . matmul ( L_S ) ) grad_x . append ( grad_u . matmul ( H_S . matmul ( L_S . t ( ) ) ) ) grad_lbda . append ( grad_u . matmul ( H_S . matmul ( - sign_z [ i ] [ S [ i ] ] ) ) ) grad_x = torch . stack ( grad_x ) grad_lbda = torch . stack ( grad_lbda ) return ( grad_x , grad_lbda )
Compute the gradient of proxTV using implicit gradient.
Compute the gradient of proxTV using implicit gradient .
"""Compute the gradient of proxTV using implicit gradient.""" # XXX do clever computations # n_dim x |S| # 1 x |S|
ctx
grad_output