How shall I modify this new_cdist() function to eliminate GPU out-of-memory runtime error ?
More specifically replacing torch.cdist() with something lighter in memory footprint.
Note: The following code snippet is related to a new type of back-propagation equation and adaptive learning rate.
See here for more information on how the existing code works
I saw there are two cdist() implementations online (code 1 , code 2)
def new_cdist(p, eta):
class cdist(torch.autograd.Function):
@staticmethod
def forward(ctx, W, X):
ctx.save_for_backward(W, X)
out = -torch.cdist(W, X, p)
return out
@staticmethod
def backward(ctx, grad_output):
W, X = ctx.saved_tensors
grad_W = grad_X = None
if ctx.needs_input_grad[0]:
_temp1 = torch.unsqueeze(X, 2).expand(X.shape[0], X.shape[1], W.shape[0]).permute(1, 0, 2)
_temp2 = torch.unsqueeze(W.transpose(0, 1), 1)
_temp = torch.cdist(_temp1, _temp2, p).squeeze().transpose(0, 1)
grad_W = torch.matmul(grad_output, _temp)
# print('before norm: ', torch.norm(grad_W))
grad_W = eta * np.sqrt(grad_W.numel()) / torch.norm(grad_W) * grad_W
print('after norm: ', torch.norm(grad_W))
if ctx.needs_input_grad[1]:
_temp1 = torch.unsqueeze(W, 2).expand(W.shape[0], W.shape[1], X.shape[0]).permute(1, 0, 2)
_temp2 = torch.unsqueeze(X.transpose(0, 1), 1)
_temp = torch.cdist(_temp1, _temp2, p).squeeze().transpose(0, 1)
_temp = torch.nn.functional.hardtanh(_temp, min_val=-1., max_val=1.)
grad_X = torch.matmul(grad_output.transpose(0, 1), _temp)
return grad_W, grad_X
return cdist().apply