From 9cfb1f622db12c1d7a511817f18b7c9783cefcf1 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Tue, 21 May 2024 13:52:05 -0300 Subject: [PATCH] softmax autograd --- .../__pycache__/functions.cpython-38.pyc | Bin 9661 -> 10296 bytes norch/autograd/functions.py | 22 +++ norch/nn/__pycache__/loss.cpython-38.pyc | Bin 2620 -> 2717 bytes norch/nn/loss.py | 12 +- tests/test_autograd.py | 155 +++++++++++++++--- 5 files changed, 158 insertions(+), 31 deletions(-) diff --git a/norch/autograd/__pycache__/functions.cpython-38.pyc b/norch/autograd/__pycache__/functions.cpython-38.pyc index 66ce27e5edd74698e2f12ebd7684abb0dc013b74..71b56620fce7e8792a0fc1d67c2e5de36dc993a3 100644 GIT binary patch delta 1140 zcmZwGZ)h8390&0GJ$IK}^ZzbM^RH=5nrzElCn5|KoGYapuBGc(eG~0!+o#RiZH-(4 z>uN%<^AE!b!`}<>4-VZ*an-(9hEr_Of`XuxDMKZwAij{{3-yhH2=#l`7lH=zeeUx- zzq@<*-QD`YJHzHfe!oX?exDCcf1Op7AJ_~h8>A>wk#tHaZjV%?Ly8nqeNvI?IZ+gl zCUr=MRLCXWkQ>qKl7ZSK0fH}r~Q1vdz}t(ukQ^y%s=^t=pf(cf0KrrKl)=t$9O9+NJn@u_#RC+ zjZhEK4DT}sXo?rjRhs9!t!L?RUblMb3~yPtXt8-Ke4XeSJ|7*SC-}SQ3cbvsSO-67 zMfqkzfXHTjR^ z3-mTWpE^v}_}kPnt@BeI7w7}F(!I3CBk5)Oi2qD4(#PEB?4=KRqw@xBHqUlN1bxAm za#6mV8K6)3=gbZIwYi-Ao#;1AKN|c6&8lb!gFL$ZFkdtenDIWlT&eWWRqgU)X9sb} z-@EtoH$fxhI#n*sR4Xlab=)pZRV#ck5lM&fj#mUOq0mC+Kc^nBIvZ}RIrr$pzcPWS zDYMEWqAu!c4bNB7Nn!7)ido3{G@cW6t)|t~xKh{W_WoU=p`wT*O^Sq4cl}dK|E{gM zQmC)#QYi_%RQiye!kuhvZbB(gOO-RHZ#%mz=(g@G>_xQnxw1WZqUElXr>m#O=k1tN z;I^F&+W{wzJ2B?OB%VR2l&2wI5c+hCzR)^xEq6oy3^*R!-olc%xS_B@d%JFe$&HqnMMMu}gR%xp*G4 z(jOniYw;&e;f-8Lsra$_<^Mf9 c_2OQtE>PSk=d=5zGaHbVtYJy>QvA8}KfpW5`Tzg` diff --git a/norch/autograd/functions.py b/norch/autograd/functions.py index f9ad107..c9a393e 100644 --- a/norch/autograd/functions.py +++ b/norch/autograd/functions.py @@ -1,4 +1,5 @@ import math +import norch class AddBackward: def __init__(self, x, y): @@ -270,3 +271,24 @@ class MinBackward: return [grad_output] +class CrossEntropyLossBackward: + def __init__(self, logits, targets): + self.input = [logits, targets] + + def backward(self, gradient): + logits, targets = self.input + + if logits.ndim == 1: + softmax = norch.softmax(logits, dim=0) + grad_logits = (softmax - targets) + + elif logits.ndim == 2: + # batched + batch_size = logits.shape[0] + softmax = norch.softmax(logits, dim=1) + + grad_logits = (softmax - targets) / batch_size + + return [grad_logits, None] # targets do not have a gradient + + diff --git a/norch/nn/__pycache__/loss.cpython-38.pyc b/norch/nn/__pycache__/loss.cpython-38.pyc index df6fd5a931db4a8f313c204e140bf1b4c3b58460..b8133b99ff9fe85bc2549c254e8f67c3b96ca772 100644 GIT binary patch delta 1113 zcma)5J#5oJ6!x9%v(q?fRMIwWQz$J|7@-cRVnU^~RJ2u#K+q)vmTNn08pjT2N6;#$ zLm3zticVrGdnY6$kPt#*XJBITz{J9W#LNKiIfz8+LZW+j@Bj0??>-b?4(2y_*#JGsHo524oVXSSQNnzC)AJ+V51x)FV2oA``9qt#hFGQ7XEiXHgcot6 zJT1a%L9-Q~$kQ#$v3=jN`o%}OvV#57G(Ln?bKv#uC~`au>FJrnIBd)r2wag4LMI9} zWL@S=+>swoAq{S#$fC?8oDhTT8aWgjY&rX>O?Na` zu2cS09J4F=U>*#5o9yUy8Xb!tY<4tw1uF2PP7eux1a5}$aK0eswHt-tb~odYM62S3 zwm3+3auDz=&bLXM9?-UG^|w{=ReQ(A0kK&SM|!z@29A6JM54sIO)pUbx7Og3Xgw{y z=#%?r(U;@igI0}uk+sEbHz_UhFp8Ffn1{{XH8?KVP93?DXy~`9fvI&kqNW_-Q&=P$ zv>Kh^dq}!)s|Msbhn~JoJ9f6MsJ7?zTy`FN&dG&^Q0&n)S!L1 zNS>6G0@`If4^bEg-8+GD<}ai9d8pUZs5{eXaNvz}XNndxldwkmx_C%Iuatt2JKHmM zD-Q8{XR58B^Om>??nJgX9j?MDr7}mNaIDkQ@*igt{c5N59KPcSd20|I{stY$6^uWB Ck?(B) delta 1036 zcma)5y=xRf6yKTIot?YeTjS-!Xf8RusCZzZ60IZ}laK^7ph1NzJoauj=90Y+W;aMe zBt=q)Rl-yv{sVS`7J`kH2o}~`3zEh@jq$xb$eq}@%l!6z{oZ>!Z|2S5rvdk2Hk&r! zDSW)U{42fh=C>b#CJjjfqe04q7?MgB&@zRu8L);~Hidtfx6~oXgtbhOOp_Kwz z=r80p;abqqYZ=hvOk4S~UW@9jW*GCptNKBxT<~<@1jg9QE%nS(GNT{UsR<6I5>fs5 zp(;tj>&^Ch^bb!|Dm5<%Div3MqYDdIKcT5X%$fyn;Duq$$B>?y4#r_&WdYy{1%MzX zb(mnv(ofj)aco1`2)NrSAtik$Osv(ff3qK#u|ahbAm)*$p88R2)>>iIQ!gRl?+j_z zx;Q(8gi{E^2;HvFpxWJ-?!*0 z40esGm-++0(YIwS8DJ7yWL2!u@U1@N6Juy>z6TSR_8Q$IY8OH{vOvC~=j@r>-@H8% zKG3i2X$OeXfKwwDH?QlfVp!+IE1rj@OIiI+E60<48(Zy{NWQ3G17F&psV*ZxLd`832R0#5poq2h#*!hq@$3 zP17ark(tPxIn<$PTo6A5zCQG^jjfoW`{qCcQADE96AK%PjlHm=MQsQ B#BBfo diff --git a/norch/nn/loss.py b/norch/nn/loss.py index bfc4390..8d7b0f9 100644 --- a/norch/nn/loss.py +++ b/norch/nn/loss.py @@ -1,4 +1,5 @@ from .module import Module +from norch.autograd.functions import * import norch from abc import ABC @@ -45,8 +46,6 @@ class CrossEntropyLoss(Loss): logits = norch.softmax(input, dim=0) cost = -(logits.log() * target).sum() - - return cost else: # target -> class probabilities (one-hot encoded) @@ -55,8 +54,6 @@ class CrossEntropyLoss(Loss): logits = norch.softmax(input, dim=0) cost = -(logits.log() * target).sum() - return cost - elif input.ndim == 2: # batched @@ -70,8 +67,6 @@ class CrossEntropyLoss(Loss): logits = norch.softmax(input, dim=1) cost = -(logits.log() * target).sum() / batch_size - return cost - else: # target -> class probabilities (one-hot encoded) assert target.shape == input.shape, \ @@ -81,7 +76,10 @@ class CrossEntropyLoss(Loss): logits = norch.softmax(input, dim=1) cost = -(logits.log() * target).sum() / batch_size - return cost + if input.requires_grad: + cost.grad_fn = CrossEntropyLossBackward(input, target) + + return cost diff --git a/tests/test_autograd.py b/tests/test_autograd.py index fb124d6..3c1b16e 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -604,40 +604,147 @@ class TestTensorAutograd(unittest.TestCase): self.assertTrue(utils.compare_torch(norch_tensor_grad, torch_tensor_grad)) - def test_softmax(self): + def test_mse_loss_autograd(self): """ - Test autograd from softmax + Test the MSELoss with autograd functionality """ - 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() + loss_fn_norch = norch.nn.MSELoss() + loss_fn_torch = torch.nn.MSELoss() - norch_result.backward() - norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device) + predictions_norch = norch.Tensor([1.1, 2, 3, 4], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([4, 3, 2.1, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad - 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() + predictions_torch = torch.tensor([1.1, 2, 3, 4], requires_grad=True).to(self.device) + labels_torch = torch.tensor([4, 3, 2.1, 1]).to(self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad - torch_result.backward() - torch_tensor_grad = torch_tensor.grad + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + def test_cross_entropy_loss_autograd(self): + """ + Test the CrossEntropyLoss with autograd functionality + """ + loss_fn_norch = norch.nn.CrossEntropyLoss() + loss_fn_torch = torch.nn.CrossEntropyLoss() + + # Test case 1: Single class, single sample + predictions_norch = norch.Tensor([2.0, 1.0, 0.1], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([0]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([2.0, 1.0, 0.1], requires_grad=True).to(self.device) + labels_torch = torch.tensor(0).to(self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 2: Multiple classes, multiple samples + predictions_norch = norch.Tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([2, 1]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True).to(self.device) + labels_torch = torch.tensor([2, 1]).to(self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 3: Edge case - all predictions are zero + predictions_norch = norch.Tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([1, 2]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True).to(self.device) + labels_torch = torch.tensor([1, 2]).to(self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + # Test case 4: Batched class probabilities instead of class index + predictions_norch = norch.Tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True).to(self.device) + labels_norch = norch.Tensor([[1., 0, 0], [0, 1, 0]]).to(self.device) + loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch) + loss_norch.backward() # Backpropagate the loss + grad_norch = predictions_norch.grad + + predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True).to(self.device) + labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]).to(self.device) + loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch) + loss_torch_expected.backward() # Backpropagate the loss + grad_torch_expected = predictions_torch.grad + + # Convert norch gradient to torch tensor for comparison + grad_norch_torch = utils.to_torch(grad_norch).to(self.device) + + self.assertTrue(utils.compare_torch(grad_norch_torch, grad_torch_expected)) + + + # implement grad pure softmax --> 0 + # 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)) + # 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_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) + # 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_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)) + # 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):