diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 141689f..4b4c091 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/tensor.py b/norch/tensor.py index 49df15d..a4e8569 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -153,6 +153,17 @@ class Tensor: result_data.grad_fn = ReshapeBackward(self) return result_data + + def unsqueeze(self, dim): + # Ensure the dimension is valid + if dim < 0 or dim > self.ndim: + raise ValueError("Dimension out of range (expected to be in range of [0, {0}], but got {1})".format(self.ndim, dim)) + + # Create the new shape with an extra dimension of size 1 + new_shape = self.shape[:dim] + [1] + self.shape[dim:] + + return self.reshape(new_shape) + def to(self, device): self.device = device diff --git a/test.py b/test.py index 54b67fb..5622e3d 100644 --- a/test.py +++ b/test.py @@ -6,37 +6,6 @@ import norch.optim as optim import random random.seed(1) -"""one_hot_target = norch.one_hot_encode(norch.Tensor([5]), num_classes=10) -print(one_hot_target)""" - -logits = norch.Tensor([[2.0, 1.0, 0.1, 0.1], [2.0, 1.0, 0.1, 0.1], [2.0, 1.0, 0.1, 0.1]], requires_grad=True) - -# One-hot encoded target with shape (batch_size, num_classes) -one_hot_target = norch.Tensor([0, 1, 1]) - -criterion = nn.CrossEntropyLoss() - -loss = criterion(logits, one_hot_target) -print(loss) - - -"""a = norch.Tensor([[[4.186502456665039]]]) -b = norch.Tensor([[[2.0, 2.0,],[-1.0, -1.0,]],[[1.0, 2.0,],[3.0, 3.0,]]]) - -print(b) -print(a.shape) -print(b-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) -""" - -""" to_tensor = lambda x: norch.Tensor(x) reshape = lambda x: x.reshape([-1, 784]) @@ -53,9 +22,9 @@ train_loader = Dataloader(train_data, batch_size = BATCH_SIZE) class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() - self.fc1 = nn.Linear(784, 10) + self.fc1 = nn.Linear(784, 5) self.sigmoid = nn.Sigmoid() - self.fc2 = nn.Linear(10, 1) + self.fc2 = nn.Linear(5, 10) def forward(self, x): out = self.fc1(x) @@ -68,7 +37,7 @@ device = "cpu" epochs = 10 model = MyModel().to(device) -criterion = nn.MSELoss() +criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.001) loss_list = [] @@ -77,13 +46,13 @@ for epoch in range(epochs): x, target = batch x = x.T - target = target.T + target = target x = x.to(device) target = target.to(device) outputs = model(x) - + print(outputs.shape, target.shape) loss = criterion(outputs, target) optimizer.zero_grad() @@ -103,4 +72,4 @@ for epoch in range(epochs): print('\n\n') print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}') - loss_list.append(loss[0])""" \ No newline at end of file + loss_list.append(loss[0]) \ No newline at end of file diff --git a/tests/test_autograd.py b/tests/test_autograd.py index 3c1b16e..c5a95f7 100644 --- a/tests/test_autograd.py +++ b/tests/test_autograd.py @@ -824,7 +824,74 @@ class TestTensorAutograd(unittest.TestCase): self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul1, torch_tensor_grad_reshape_matmul1)) self.assertTrue(utils.compare_torch(norch_tensor_grad_reshape_matmul2, torch_tensor_grad_reshape_matmul2)) + def test_unsqueeze(self): + """ + Test autograd from unsqueezing a tensor: tensor.unsqueeze(dim) + """ + + # Unsqueeze at dim=0 + norch_tensor_unsqueeze = norch.Tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_0 = norch_tensor_unsqueeze.unsqueeze(0).sum() + norch_result_unsqueeze_0.backward() + norch_tensor_grad_unsqueeze_0 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + torch_result_unsqueeze_0 = torch_tensor_unsqueeze.unsqueeze(0).sum() + torch_result_unsqueeze_0.backward() + torch_tensor_grad_unsqueeze_0 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_0, torch_tensor_grad_unsqueeze_0)) + + # Unsqueeze at dim=1 + norch_tensor_unsqueeze = norch.Tensor([[1., 2.], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_1 = norch_tensor_unsqueeze.unsqueeze(1).sum() + norch_result_unsqueeze_1.backward() + norch_tensor_grad_unsqueeze_1 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + + torch_tensor_unsqueeze = torch.tensor([[1., 2.], [3, 4]], requires_grad=True).to(self.device) + torch_result_unsqueeze_1 = torch_tensor_unsqueeze.unsqueeze(1).sum() + torch_result_unsqueeze_1.backward() + torch_tensor_grad_unsqueeze_1 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_1, torch_tensor_grad_unsqueeze_1)) + + # Unsqueeze at dim=2 + norch_tensor_unsqueeze = norch.Tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + norch_result_unsqueeze_2 = norch_tensor_unsqueeze.unsqueeze(2).sum() + norch_result_unsqueeze_2.backward() + norch_tensor_grad_unsqueeze_2 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device) + + torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device) + torch_result_unsqueeze_2 = torch_tensor_unsqueeze.unsqueeze(2).sum() + torch_result_unsqueeze_2.backward() + torch_tensor_grad_unsqueeze_2 = torch_tensor_unsqueeze.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_2, torch_tensor_grad_unsqueeze_2)) + + def test_unsqueeze_then_matmul(self): + """ + Test autograd from unsqueezing a tensor then performing matrix multiplication: matmul(tensor1.unsqueeze(dim), tensor2) + """ + norch_tensor1 = norch.Tensor([[1, 2], [3, 4]], requires_grad=True).to(self.device) + norch_tensor2 = norch.Tensor([[1, 2], [3, 4]], requires_grad=True).to(self.device) + + # Unsqueeze at dim=0 then matmul + norch_result_unsqueeze_matmul = (norch_tensor1 @ norch_tensor2.unsqueeze(0)).sum() + norch_result_unsqueeze_matmul.backward() + norch_tensor_grad_unsqueeze_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device) + norch_tensor_grad_unsqueeze_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device) + + torch_tensor1 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True).to(self.device) + torch_tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True).to(self.device) + + torch_result_unsqueeze_matmul = (torch_tensor1 @ torch_tensor2.unsqueeze(0)).sum() + torch_result_unsqueeze_matmul.backward() + torch_tensor_grad_unsqueeze_matmul1 = torch_tensor1.grad + torch_tensor_grad_unsqueeze_matmul2 = torch_tensor2.grad + + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_matmul1, torch_tensor_grad_unsqueeze_matmul1)) + self.assertTrue(utils.compare_torch(norch_tensor_grad_unsqueeze_matmul2, torch_tensor_grad_unsqueeze_matmul2)) + def test_T_then_matmul(self): """ Test autograd from transposing a tensor then performing matrix multiplication: matmul(tensor.T, tensor) diff --git a/tests/test_operations.py b/tests/test_operations.py index 961cb17..e141066 100644 --- a/tests/test_operations.py +++ b/tests/test_operations.py @@ -207,6 +207,31 @@ class TestTensorOperations(unittest.TestCase): self.assertTrue(utils.compare_torch(torch_result, torch_expected)) + def test_unsqueeze(self): + """ + Test unsqueeze operation on a tensor + """ + norch_tensor = norch.Tensor([[1, 2], [3, 4]]).to(self.device) + + # Unsqueeze at dim=0 + norch_unsqueeze_0 = norch_tensor.unsqueeze(0) + torch_unsqueeze_0 = utils.to_torch(norch_unsqueeze_0).to(self.device) + torch_tensor = torch.tensor([[1, 2], [3, 4]]).to(self.device) + torch_expected_0 = torch_tensor.unsqueeze(0) + self.assertTrue(utils.compare_torch(torch_unsqueeze_0, torch_expected_0)) + + # Unsqueeze at dim=1 + norch_unsqueeze_1 = norch_tensor.unsqueeze(1) + torch_unsqueeze_1 = utils.to_torch(norch_unsqueeze_1).to(self.device) + torch_expected_1 = torch_tensor.unsqueeze(1) + self.assertTrue(utils.compare_torch(torch_unsqueeze_1, torch_expected_1)) + + # Unsqueeze at dim=2 + norch_unsqueeze_2 = norch_tensor.unsqueeze(2) + torch_unsqueeze_2 = utils.to_torch(norch_unsqueeze_2).to(self.device) + torch_expected_2 = torch_tensor.unsqueeze(2) + self.assertTrue(utils.compare_torch(torch_unsqueeze_2, torch_expected_2)) + def test_transpose(self): """ Test transposition of a tensor: tensor.transpose(dim1, dim2)