min max broadcasted autograd and small fix on add broadcasted

This commit is contained in:
lucasdelimanogueira 2024-05-20 18:44:43 -03:00
parent b324b7e1fa
commit cd57c39ff0
7 changed files with 119 additions and 26 deletions

View file

@ -226,15 +226,10 @@ class MaxBackward:
grad_output_shape.insert(self.axis, 1)
grad_output = gradient.reshape(grad_output_shape)
grad_output = grad_output + self.input[0].zeros_like()
print(self.input[0])
max_value = self.input[0].max()
print(max_value)
max_values = self.input[0].max(axis=self.axis, keepdim=True)
print('\n\n', max_values, '@@@')
mask = self.input[0].equal(max_value)
mask = self.input[0].equal(max_values)
grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
grad_output = (grad_output * mask)
return [grad_output]
@ -256,7 +251,19 @@ class MinBackward:
grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
else:
pass
if not self.keepdim:
# Remove dimensions of size 1 from the gradient tensor.
input_shape = [s for i, s in enumerate(input_shape) if i != self.axis]
# Broadcast the gradient to the input shape along the specified axis.
grad_output_shape = list(input_shape)
grad_output_shape.insert(self.axis, 1)
grad_output = gradient.reshape(grad_output_shape)
grad_output = grad_output + self.input[0].zeros_like()
max_values = self.input[0].min(axis=self.axis, keepdim=True)
mask = self.input[0].equal(max_values)
grad_output = (grad_output * mask)
return [grad_output]

View file

@ -270,6 +270,14 @@ class Tensor:
if needs_broadcasting:
# Call add_broadcasted_tensor if broadcasting is needed
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
Tensor._C.add_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
@ -361,6 +369,12 @@ class Tensor:
broadcasted_shape_sub, needs_broadcasting = broadcast_shape(self.shape, other.shape)
if needs_broadcasting:
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
# Call add_broadcasted_tensor if broadcasting is needed
Tensor._C.sub_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.sub_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
@ -678,6 +692,12 @@ class Tensor:
broadcasted_shape_add, needs_broadcasting = broadcast_shape(self.shape, other.shape)
if needs_broadcasting:
if other.ndim == self.ndim - 1:
other = other.reshape([1] + other.shape)
elif self.ndim == other.ndim - 1:
self = self.reshape([1] + self.shape)
# Call equal_broadcasted_tensor if broadcasting is needed
Tensor._C.equal_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.equal_broadcasted_tensor.restype = ctypes.POINTER(CTensor)

18
test.py
View file

@ -6,12 +6,20 @@ import norch.optim as optim
import random
random.seed(1)
a = norch.Tensor([1, 2, 3])
b = 3
for i in range(a.numel):
print(a.tensor.contents.data[i])
torch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)#.to(self.device)
torch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]])
torch_tensor3 = torch_tensor.max(axis=2, keepdim=True)
print(torch_tensor3)
print(torch_tensor.equal(torch_tensor3))
print(a)
"""print(torch_tensor.shape)
print('\n\n')
torch_tensor2 = torch_tensor.max(axis=1)
print(torch_tensor2.shape)
print('\n\n')
c = torch_tensor + torch_tensor2
print(c)
"""
"""

View file

@ -115,24 +115,65 @@ class TestTensorAutograd(unittest.TestCase):
self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
norch_max_axis = norch_tensor.max(axis=2)
norch_result = norch_max_axis.sum()
norch_result.backward()
norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
torch_max_axis, _ = torch_tensor.max(axis=2)
torch_result = torch_max_axis.sum()
torch_result.backward()
torch_tensor_grad = torch_tensor.grad
self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
print(norch_tensor_grad, '\n\n', torch_tensor_grad)
## evaluate case with some repeated values and axis 2
#def test_max_axis(self):
# """
# Test autograd from max specifying axis
# """
#
# norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
# norch_max_axis = norch_tensor.max(axis=2)
# norch_result = norch_max_axis.sum()
#
# norch_result.backward()
# norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
#
# torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
#
# torch_max_axis, _ = torch_tensor.max(axis=2)
# torch_result = torch_max_axis.sum()
# torch_result.backward()
# torch_tensor_grad = torch_tensor.grad
# self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
#def test_min_axis(self):
# """
# Test autograd from max specifying axis
# """
#
# norch_tensor = norch.Tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
# norch_min_axis = norch_tensor.min(axis=2)
# norch_result = norch_min_axis.sum()
#
# norch_result.backward()
# norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
#
# torch_tensor = torch.tensor([[[10, 10], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
#
# torch_min_axis, _ = torch_tensor.min(axis=2)
# torch_result = torch_min_axis.sum()
# torch_result.backward()
# torch_tensor_grad = torch_tensor.grad
# self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
def test_min(self):
@ -166,23 +207,23 @@ class TestTensorAutograd(unittest.TestCase):
torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
torch_min = torch_tensor.min(axis=1)
torch_min, _ = torch_tensor.min(axis=1)
torch_result = torch_min.sum()
torch_result.backward()
torch_tensor_grad = torch_tensor.grad
self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_tensor = norch.Tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
norch_min = norch_tensor.min(axis=2)
norch_result = norch_min.sum()
norch_result.backward()
norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
torch_min = torch_tensor.min(axis=2)
torch_min, _ = torch_tensor.min(axis=2)
torch_result = torch_min.sum()
torch_result.backward()
torch_tensor_grad = torch_tensor.grad

View file

@ -52,6 +52,17 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3)
norch_tensor2 = norch.Tensor([[10, 10], [5, 6]]).to(self.device) # Shape (3)
norch_result = norch_tensor1 + norch_tensor2
torch_result = utils.to_torch(norch_result).to(self.device)
torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device) # Shape (1, 2, 3)
torch_tensor2 = torch.tensor([[[10, 10], [5, 6]]]).to(self.device) # Shape (3)
torch_expected = torch_tensor1 + torch_tensor2
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
# reversed order broadcasting
norch_tensor1 = norch.Tensor([[0, 2]]).to(self.device)
norch_tensor2 = norch.Tensor([[3, 4], [5, -1]]).to(self.device)
@ -274,8 +285,6 @@ class TestTensorOperations(unittest.TestCase):
torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
torch_expected = torch.max(torch_tensor)
print(torch_result, torch_expected)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_max_axis(self):
@ -316,8 +325,6 @@ class TestTensorOperations(unittest.TestCase):
torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
torch_expected = torch.min(torch_tensor)
print(torch_result, torch_expected)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_min_axis(self):
@ -561,7 +568,17 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
norch_tensor1 = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
norch_tensor2 = norch.Tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device)
norch_result = norch_tensor1.equal(norch_tensor2)
torch_result = utils.to_torch(norch_result).to(self.device)
torch_tensor1 = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]]).to(self.device)
torch_tensor2 = torch.tensor([[[10.0,], [-4.0,]],[[6.0,], [8.0,]]]).to(self.device)
torch_expected = (torch_tensor1 == torch_tensor2).float()
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_zeros_like(self):
"""