aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #52 from lucasdelimanogueira/tmp

Browse source 
Tmp
This commit is contained in:
lucasdelimanogueira 2024-05-16 16:38:30 -03:00 • committed by GitHub
commit 561d3a7e01
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 266 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
build/cpu.o Normal file
View file

Binary file not shown.

BIN
build/cuda.cu.o Normal file
View file

Binary file not shown.

BIN
build/tensor.o Normal file
View file

Binary file not shown.

BIN
norch/__pycache__/tensor.cpython-38.pyc
View file

Binary file not shown.

BIN
norch/autograd/__pycache__/functions.cpython-38.pyc
View file

Binary file not shown.

6
norch/autograd/functions.py
View file

@ -7,6 +7,12 @@ class AddBackward:
def backward(self, gradient):
return [gradient, gradient]
class AddBroadcastedBackward:
pass
class SubBackward:
def __init__(self, x, y):
self.input = [x, y]

88
norch/csrc/cpu.cpp
View file

@ -11,6 +11,46 @@ void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
}
}
void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape) {
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
// Calculate strides for broadcasting
int* strides1 = (int*)malloc(max_ndim * sizeof(int));
int* strides2 = (int*)malloc(max_ndim * sizeof(int));
if (strides1 == NULL || strides2 == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
int stride1 = 1, stride2 = 1;
for (int i = max_ndim - 1; i >= 0; i--) {
int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - max_ndim + i] : 1;
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - max_ndim + i] : 1;
strides1[i] = dim1 == broadcasted_shape[i] ? stride1 : 0;
strides2[i] = dim2 == broadcasted_shape[i] ? stride2 : 0;
stride1 *= broadcasted_shape[i];
stride2 *= broadcasted_shape[i];
}
// Perform element-wise addition with broadcasting
for (int i = 0; i < tensor1->size; i++) {
int index1 = 0, index2 = 0;
int linear_index = i;
for (int j = max_ndim - 1; j >= 0; j--) {
int pos = linear_index % broadcasted_shape[j];
linear_index /= broadcasted_shape[j];
if (strides1[j] != 0) index1 += pos * strides1[j];
if (strides2[j] != 0) index2 += pos * strides2[j];
}
result_data[i] = tensor1->data[index1] + tensor2->data[index2];
}
// Free strides
free(strides1);
free(strides2);
}
void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
for (int i = 0; i < tensor1->size; i++) {
@ -126,16 +166,52 @@ void log_tensor_cpu(Tensor* tensor, float* result_data) {
}
}
void sum_tensor_cpu(Tensor* tensor, float* result_data) {
float sum = 0.0;
void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis) {
if (axis == -1) {
// Sum over all elements
float sum = 0.0;
for (int i = 0; i < tensor->size; i++) {
sum += tensor->data[i];
}
*result_data = sum;
} else {
if (axis < 0 || axis >= tensor->ndim) {
printf("Invalid axis");
return;
}
int result_shape[tensor->ndim - 1];
int result_size = 1;
int axis_stride = tensor->strides[axis];
for (int i = 0; i < tensor->size; i++) {
sum += tensor->data[i];
int idx = 0;
for (int i = 0; i < tensor->ndim; i++) {
if (i != axis) {
result_shape[idx++] = tensor->shape[i];
result_size *= tensor->shape[i];
}
}
memset(result_data, 0, result_size * sizeof(float));
for (int i = 0; i < tensor->shape[axis]; i++) {
for (int j = 0; j < result_size; j++) {
int index = 0;
int remainder = j;
for (int k = tensor->ndim - 2; k >= 0; k--) {
index += (remainder % result_shape[k]) * tensor->strides[k < axis ? k : k + 1];
remainder /= result_shape[k];
}
result_data[j] += tensor->data[index + i * axis_stride];
}
}
}
*result_data = sum;
}
void ones_like_tensor_cpu(Tensor* tensor, float* result_data) {
for (int i = 0; i < tensor->size; i++) {

3
norch/csrc/cpu.h
View file

@ -4,7 +4,8 @@
#include "tensor.h"
void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void sum_tensor_cpu(Tensor* tensor1, float* result_data);
void add_broadcasted_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape);
void sum_tensor_cpu(Tensor* tensor, float* result_data, int axis);
void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data);
void scalar_div_tensor_cpu(float scalar, Tensor* tensor, float* result_data);

66
norch/csrc/tensor.cpp
View file

@ -124,7 +124,44 @@ extern "C" {
}
}
Tensor* sum_tensor(Tensor* tensor) {
Tensor* add_broadcasted_tensor(Tensor* tensor1, Tensor* tensor2) {
if (strcmp(tensor1->device, tensor2->device) != 0) {
fprintf(stderr, "Tensors must be on the same device: %s and %s\n", tensor1->device, tensor2->device);
exit(1);
}
int max_ndim = tensor1->ndim > tensor2->ndim ? tensor1->ndim : tensor2->ndim;
// Determine the broadcasted shape
int* broadcasted_shape = (int*)malloc(max_ndim * sizeof(int));
if (broadcasted_shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
for (int i = 0; i < max_ndim; i++) {
int dim1 = i < tensor1->ndim ? tensor1->shape[tensor1->ndim - 1 - i] : 1;
int dim2 = i < tensor2->ndim ? tensor2->shape[tensor2->ndim - 1 - i] : 1;
if (dim1 != dim2 && dim1 != 1 && dim2 != 1) {
fprintf(stderr, "Shapes are not compatible for broadcasting\n");
exit(1);
}
broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
}
// Allocate memory for result tensor
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
Tensor* sum_tensor(Tensor* tensor, int axis) {
char* device = (char*)malloc(strlen(tensor->device) + 1);
if (device != NULL) {
@ -133,14 +170,25 @@ extern "C" {
fprintf(stderr, "Memory allocation failed\n");
exit(-1);
}
int ndim = 1;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
int ndim;
int* shape;
if (axis == -1) {
shape = (int*) malloc(sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
shape[0] = 1;
ndim = 1;
} else {
shape = (int*) malloc((tensor->ndim - 1) * sizeof(int));
for (int i = 0, j = 0; i < tensor->ndim; ++i) {
if (i != axis) {
shape[j++] = tensor->shape[i];
}
}
ndim = tensor->ndim - 1;
}
shape[0] = 1;
if (strcmp(tensor->device, "cuda") == 0) {
@ -155,7 +203,7 @@ extern "C" {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
sum_tensor_cpu(tensor, result_data);
sum_tensor_cpu(tensor, result_data, axis);
return create_tensor(result_data, shape, ndim, device);
}
}

2
norch/csrc/tensor.h
View file

@ -16,7 +16,7 @@ extern "C" {
Tensor* create_tensor(float* data, int* shape, int ndim, char* device);
float get_item(Tensor* tensor, int* indices);
Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2);
Tensor* sum_tensor(Tensor* tensor);
Tensor* sum_tensor(Tensor* tensor, int axis);
Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2);
Tensor* elementwise_mul_tensor(Tensor* tensor1, Tensor* tensor2);
Tensor* scalar_mul_tensor(Tensor* tensor, float scalar);

BIN
norch/libtensor.so
View file

Binary file not shown.

87
norch/tensor.py
View file

@ -223,27 +223,65 @@ class Tensor:
def __add__(self, other):
if isinstance(other, (int, float)):
other = other * self.ones_like()
# Function to determine if broadcasting is needed and get the broadcasted shape
def broadcast_shape(shape1, shape2):
if shape1 == shape2:
return shape1, False
max_len = max(len(shape1), len(shape2))
shape1 = [1] * (max_len - len(shape1)) + shape1
shape2 = [1] * (max_len - len(shape2)) + shape2
broadcasted_shape = []
for dim1, dim2 in zip(shape1, shape2):
if dim1 != dim2 and dim1 != 1 and dim2 != 1:
raise ValueError("Shapes are not compatible for broadcasting")
broadcasted_shape.append(max(dim1, dim2))
return broadcasted_shape, True
broadcasted_shape, needs_broadcasting = broadcast_shape(self.shape, other.shape)
if needs_broadcasting:
# Call add_broadcasted_tensor if broadcasting is needed
Tensor._C.add_broadcasted_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_broadcasted_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.add_broadcasted_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = broadcasted_shape
result_data.ndim = len(broadcasted_shape)
result_data.device = self.device
result_data.numel = self.numel # Update this to calculate the correct number of elements if broadcasting
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBroadcastedBackward(self, other)
if self.shape != other.shape:
raise ValueError("Tensors must have the same shape for addition")
Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor)
else:
# Call add_tensor if shapes are identical
Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.add_tensor(self.tensor, other.tensor)
result_tensor_ptr = Tensor._C.add_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
result_data.numel = self.numel
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBackward(self, other)
result_data.device = self.device
result_data.numel = self.numel # Update this to calculate the correct number of elements if broadcasting
result_data.requires_grad = self.requires_grad or other.requires_grad
if result_data.requires_grad:
result_data.grad_fn = AddBackward(self, other)
return result_data
def __radd__(self, other):
if isinstance(other, (int, float)):
@ -555,24 +593,33 @@ class Tensor:
return result_data
def sum(self):
Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor)]
def sum(self, axis=-1):
Tensor._C.sum_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.c_int]
Tensor._C.sum_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.sum_tensor(self.tensor)
result_tensor_ptr = Tensor._C.sum_tensor(self.tensor, axis)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = [1]
result_data.ndim = 1
if axis == -1:
result_data.shape = [1]
result_data.ndim = 1
else:
result_data.shape = self.shape[:axis] + self.shape[axis+1:]
result_data.ndim = len(result_data.shape)
result_data.device = self.device
result_data.numel = 1
for s in result_data.shape:
result_data.numel *= s
result_data.requires_grad = self.requires_grad
if result_data.requires_grad:
result_data.grad_fn = SumBackward(self)
return result_data
def sin(self):
Tensor._C.sin_tensor.argtypes = [ctypes.POINTER(CTensor)]

20
tests/test_autograd.py
View file

@ -31,6 +31,26 @@ 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))
def test_broadcasting_addition_autograd(self):
"""
Test autograd for broadcasting addition: tensor1 + tensor2
"""
norch_tensor1 = norch.Tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device) # Shape (1, 2, 3)
norch_tensor2 = norch.Tensor([1.5, -1, 0], requires_grad=True).to(self.device) # Shape (3)
norch_result = (norch_tensor1 + norch_tensor2).sum()
norch_result.backward()
norch_tensor1_grad = utils.to_torch(norch_tensor1.grad)
norch_tensor2_grad = utils.to_torch(norch_tensor2.grad)
torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device) # Shape (1, 2, 3)
torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device) # Shape (3)
torch_result = (torch_tensor1 + torch_tensor2).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_subtraction(self):
"""

31
tests/test_operations.py
View file

@ -35,6 +35,22 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_broadcasting_addition(self):
"""
Test addition of two tensors with broadcasting: tensor1 + tensor2
"""
norch_tensor1 = norch.Tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3)
norch_tensor2 = norch.Tensor([1, 1, 1]).to(self.device) # Shape (3)
norch_result = norch_tensor1 + norch_tensor2
torch_result = utils.to_torch(norch_result)
torch_tensor1 = torch.tensor([[[1, 2, 3], [4, 5, 6]]]).to(self.device) # Shape (1, 2, 3)
torch_tensor2 = torch.tensor([1, 1, 1]).to(self.device) # Shape (3)
torch_expected = torch_tensor1 + torch_tensor2 # Broadcasting addition
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_subtraction(self):
"""
Test subtraction of two tensors: tensor1 - tensor2
@ -176,6 +192,21 @@ class TestTensorOperations(unittest.TestCase):
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_sum_axis_no_keepdims(self):
"""
Test summation of a tensor along a specific axis without keeping the dimensions
"""
norch_tensor = norch.Tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
norch_result = norch_tensor.sum(axis=1)
torch_result = utils.to_torch(norch_result)
torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
torch_expected = torch.sum(torch_tensor, dim=1)
print(torch_result)
print("\n\n", torch_expected)
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
def test_transpose_T(self):
"""
Test transposition of a tensor: tensor.T