min max autograd

This commit is contained in:
lucasdelimanogueira 2024-05-20 15:24:55 -03:00
parent cc483417dc
commit 098ede5a6b
7 changed files with 158 additions and 11 deletions

View file

@ -206,8 +206,34 @@ class MaxBackward:
self.keepdim = keepdim
def backward(self, gradient):
pass
input_shape = self.input[0].shape.copy()
if self.axis == -1:
max_value = self.input[0].max()
mask = self.input[0].equal(max_value)
grad_output = float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()
grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
else:
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_value = self.input[0].max()
mask = self.input[0].equal(max_value)
grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
return [grad_output]
class MinBackward:
def __init__(self, x, axis=None, keepdim=False):
self.input = [x]
@ -215,6 +241,18 @@ class MinBackward:
self.keepdim = keepdim
def backward(self, gradient):
pass
input_shape = self.input[0].shape.copy()
if self.axis == -1:
min_value = self.input[0].min()
mask = self.input[0].equal(min_value)
grad_output = float(gradient.tensor.contents.data[0]) * self.input[0].ones_like()
grad_output = (grad_output * mask) / mask.sum().tensor.contents.data[0]
else:
pass
return [grad_output]

View file

@ -643,8 +643,19 @@ class Tensor:
def equal(self, other):
if isinstance(other, Tensor) and other.numel == 1:
# other is a single value tensor
other = self.zeros_like() + other
return self.equal(other)
if not isinstance(other, Tensor):
return False
# other is a single value
if isinstance(other, (int, float)):
other = self.zeros_like() + other
return self.equal(other)
else:
return False
if self.shape != other.shape:
return False

View file

@ -7,8 +7,11 @@ import random
random.seed(1)
a = norch.Tensor([1, 2, 3])
b = norch.Tensor([1, 2, 4])
print(a == b)
b = 3
for i in range(a.numel):
print(a.tensor.contents.data[i])
print(a)
"""

View file

@ -56,39 +56,134 @@ class TestTensorAutograd(unittest.TestCase):
self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
norch_tensor1 = norch.Tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
norch_tensor2 = norch.Tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
norch_result = (norch_tensor1 + norch_tensor2).sum(axis=1).sum()
norch_result.backward()
norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
torch_result = (torch_tensor1 + torch_tensor2).sum(axis=1).sum()
torch_result.backward()
torch_tensor1_grad = torch_tensor1.grad
torch_tensor2_grad = torch_tensor2.grad
self.assertTrue(utils.compare_torch(norch_tensor1_grad, torch_tensor1_grad))
self.assertTrue(utils.compare_torch(norch_tensor2_grad, torch_tensor2_grad))
def test_max(self):
"""
Test autograd from max
"""
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_result = norch_tensor.max()
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_result = torch_tensor.max()
torch_result.backward()
torch_tensor_grad = torch_tensor.grad
self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad))
def test_max_axis(self):
"""
Test autograd from max specifying axis
"""
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_result = norch_tensor.max(axis=1).sum()
norch_max_axis = norch_tensor.max(axis=1)
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_max_axis, _ = torch_tensor.max(axis=1)
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))
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [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_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)
def test_min(self):
"""
Test autograd from min
"""
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_result = norch_tensor.min()
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_result = torch_tensor.max(axis=1).sum()
torch_result = torch_tensor.min()
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 min specifying axis
"""
norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
norch_result = norch_tensor.min(axis=1).sum()
norch_min = norch_tensor.min(axis=1)
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_result = torch_tensor.min(axis=1).sum()
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_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_min = torch_tensor.min(axis=2)
torch_result = torch_min.sum()
torch_result.backward()
torch_tensor_grad = torch_tensor.grad

View file

@ -533,7 +533,7 @@ class TestTensorOperations(unittest.TestCase):
def test_equal(self):
"""
Test equal two tensors: tensor1 == tensor2
Test equal two tensors: tensor1.equal(tensor2)
"""
norch_tensor1 = norch.Tensor([[[1, 2], [3, -4]], [[5, 1], [7, 8]]]).to(self.device)
norch_tensor2 = norch.Tensor([[[1, 1], [1, 1]], [[1, 1], [1, 1]]]).to(self.device)