diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc index 0b21229..66ce27e 100644 Binary files a/norch/autograd/__pycache__/functions.cpython-38.pyc and b/norch/autograd/__pycache__/functions.cpython-38.pyc differ diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py index e64841a..f9ad107 100644 --- a/norch/autograd/functions.py +++ b/norch/autograd/functions.py @@ -137,13 +137,14 @@ class SumBackward: # If axis is None, sum reduces the tensor to a scalar. grad_output = float(gradient.tensor.contents.data[0]) * self.input[0].ones_like() 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] - + + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + # 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() @@ -216,14 +217,15 @@ class MaxBackward: 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] - + + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + # 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].max(axis=self.axis, keepdim=True) @@ -251,13 +253,13 @@ class MinBackward: 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] - + if self.keepdim: + input_shape = input_shape[:self.axis] + [1] + input_shape[self.axis+1:] + else: + input_shape = input_shape[:self.axis] + input_shape[self.axis+1:] + # 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) diff --git a/test.py b/test.py index 6bd8b2a..d61d329 100644 --- a/test.py +++ b/test.py @@ -6,10 +6,12 @@ import norch.optim as optim import random random.seed(1) -torch_tensor = norch.Tensor([[[2, 2], [-1, -1]], [[1., 2], [3, 3]]], requires_grad=True)#.to(self.device) -torch_tensor = norch.Tensor([[[2, 2], [2, 2]], [[2, 2.], [2, 2]]]) -b = norch.nn.functional.softmax(torch_tensor) -print(b) +torch_tensor = norch.Tensor([[[1, 5], [2, -1]], [[5, 2.], [2, 2]]], requires_grad=True)#.to(self.device) +b = norch.nn.functional.softmax(torch_tensor, dim=0) +soma = b.sum() +print(soma) +soma.backward() +print(torch_tensor.grad) """a = norch.Tensor([[[4.186502456665039]]]) b = norch.Tensor([[[2.0, 2.0,],[-1.0, -1.0,]],[[1.0, 2.0,],[3.0, 3.0,]]]) diff --git a/tests/test_autograd.py b/tests/test_autograd.py index eb27544..fb124d6 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -583,6 +583,62 @@ class TestTensorAutograd(unittest.TestCase): torch_expected_cos_tensor_grad = torch_cos_tensor.grad self.assertTrue(utils.compare_torch(torch_result_cos_tensor_grad, torch_expected_cos_tensor_grad)) + + def test_sigmoid(self): + """ + Test autograd from sigmoid + """ + norch_tensor = norch.Tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_sigmoid = norch.sigmoid(norch_tensor) + norch_result = norch_sigmoid.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_sigmoid = torch.sigmoid(torch_tensor) + torch_result = torch_sigmoid.sum() + + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + + def test_softmax(self): + """ + Test autograd from softmax + """ + norch_tensor = norch.Tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + norch_softmax = norch.softmax(norch_tensor, dim=1) + norch_result = norch_softmax.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[-5, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device) + torch_softmax = torch.softmax(torch_tensor, dim=1) + torch_result = torch_softmax.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, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + norch_softmax = norch.softmax(norch_tensor, dim=2) + norch_result = norch_softmax.sum() + + norch_result.backward() + norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + + torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device) + + torch_softmax = torch.softmax(torch_tensor, dim=2) + torch_result = torch_softmax.sum() + torch_result.backward() + torch_tensor_grad = torch_tensor.grad + self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) + def test_reshape(self): """