Save memory when compute shortest distance

[Tiamo666]

[Tiamo666]

@michuanhaohao Thanks for your great work. But I had a suggestion that when compute the shortest distance will cost O(MN) extra space, where M, N is the size of the dist_mat. We can reduce the extra space from O(MN) to O(N). I implemented it as following:

def shortest_path(dist_mat):
    """
    Args:
    dist_mat: pytorch Variable, available shape:
      1) [m, n]
      2) [m, n, N], N is batch size
      3) [m, n, *], * can be arbitrary additional dimensions
      Returns:
    dist: three cases corresponding to `dist_mat`:
      1) scalar
      2) pytorch Tensor, with shape [N]
      3) pytorch Tensor, with shape [*]
    """
    if dist_mat.dim() == 3:
        M, N, B = dist_mat.size()
        helper = torch.zeros(N, B)
    elif dist_mat.dim() == 2:
        M, N = dist.size()
        helper = torch.zeros(N)
    for i in range(M):
        for j in range(N):
            if i == 0 and j == 0:
                helper[j] = dist_mat[i][j]
            elif i == 0:
                helper[j] = helper[j-1] + dist_mat[i][j]
            elif j == 0:
                helper[j] = helper[j] + dist_mat[i][j]
            else:
                helper[j] = torch.min(helper[j-1], helper[j]) + dist_mat[i][j]
    return helper[-1]

I think we can even do better to reduce the extra space to O(min(N, M))

[michuanhaohao]

Thank you for you job! If results can be reproduced by your new implementation, you can make a pull request and become a contributor of this project.

[Tiamo666]

OK, but I was a little busy with other things recently, I'll make a pull request if I can reimplement the result.