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

fix memory usage

Browse source
This commit is contained in:
lucasdelimanogueira 2024-06-05 21:04:09 -03:00
parent 7ac5fbcc47
commit 70cbd45ae7
8 changed files with 148 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
build/cuda.cu.o
View file

Binary file not shown.

BIN
build/tensor.o
View file

Binary file not shown.

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

Binary file not shown.

62
norch/csrc/tensor.cpp
View file

@ -49,37 +49,30 @@ extern "C" {
void delete_tensor(Tensor* tensor) {
if (tensor != NULL) {
if (tensor->shape != NULL) {
free(tensor->shape);
tensor->shape = NULL;
}
if (tensor->data != NULL) {
if (strcmp(tensor->device, "cpu") == 0) {
free(tensor->data);
} else {
free_cuda(tensor->data);
}
tensor->data = NULL;
}
if (tensor->device != NULL) {
free(tensor->device);
tensor->device = NULL;
}
if (tensor->strides != NULL) {
free(tensor->strides);
tensor->strides = NULL;
}
free(tensor);
tensor = NULL;
}
}
void delete_shape(Tensor* tensor) {
if (tensor->shape != NULL) {
free(tensor->shape);
tensor->shape = NULL;
}
}
void delete_data(Tensor* tensor) {
if (tensor->data != NULL) {
if (strcmp(tensor->device, "cpu") == 0) {
free(tensor->data);
} else {
free_cuda(tensor->data);
}
tensor->data = NULL;
}
}
void delete_strides(Tensor* tensor) {
// as strides always are allocated within C code, it must be handled separatedly
if (tensor->strides != NULL) {
free(tensor->strides);
tensor->strides = NULL;
@ -87,7 +80,6 @@ extern "C" {
}
void delete_device(Tensor* tensor) {
// as device always are allocated within C code, it must be handled separatedly
if (tensor->device != NULL) {
free(tensor->device);
tensor->device = NULL;
@ -114,20 +106,28 @@ extern "C" {
int device_id = 0;
char* endptr;
char* target_device_type;
long num = strtol(target_device, &endptr, 10);
if (*endptr == '\0') {
device_id = (int)num;
target_device = new char[strlen("cuda") + 1];
strcpy(target_device, "cuda");
target_device_type = new char[strlen("cuda") + 1];
strcpy(target_device_type, "cuda");
}
else {
target_device_type = new char[strlen("cuda") + 1];
strcpy(target_device_type, "cpu");
}
if ((strcmp(target_device, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
if ((strcmp(target_device_type, "cuda") == 0) && (strcmp(tensor->device, "cpu") == 0)) {
cpu_to_cuda(tensor, device_id);
}
else if ((strcmp(target_device, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {
else if ((strcmp(target_device_type, "cpu") == 0) && (strcmp(tensor->device, "cuda") == 0)) {
cuda_to_cpu(tensor);
}
free(target_device_type);
}
Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2) {

2
norch/csrc/tensor.h
View file

@ -14,6 +14,8 @@ extern "C" {
Tensor* create_tensor(float* data, int* shape, int ndim, char* device);
void delete_tensor(Tensor* tensor);
void delete_strides(Tensor* tensor);
void delete_shape(Tensor* tensor);
void delete_strides(Tensor* tensor);
void delete_device(Tensor* tensor);
float get_item(Tensor* tensor, int* indices);
Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2);

BIN
norch/libtensor.so
View file

Binary file not shown.

23
norch/tensor.py
View file

@ -83,8 +83,7 @@ class Tensor:
def __del__(self):
if hasattr(self, '_data_ctype') and self._data_ctype is not None:
# tensor created by user (ctypes) will be deleted by python garbage collector
# only strides need to be deallocated manually because it is created inside C code
Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_strides.restype = None
Tensor._C.delete_strides(self.tensor)
@ -93,8 +92,26 @@ class Tensor:
Tensor._C.delete_device.restype = None
Tensor._C.delete_device(self.tensor)
Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_tensor.restype = None
Tensor._C.delete_tensor(self.tensor)
elif self.tensor is not None:
# tensor created during operations must be deallocated
Tensor._C.delete_strides.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_strides.restype = None
Tensor._C.delete_strides(self.tensor)
Tensor._C.delete_data.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_data.restype = None
Tensor._C.delete_data(self.tensor)
Tensor._C.delete_shape.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_shape.restype = None
Tensor._C.delete_shape(self.tensor)
Tensor._C.delete_device.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_device.restype = None
Tensor._C.delete_device(self.tensor)
Tensor._C.delete_tensor.argtypes = [ctypes.POINTER(CTensor)]
Tensor._C.delete_tensor.restype = None

95
train_singlegpu.py Normal file
View file

@ -0,0 +1,95 @@
import norch
import norch.nn as nn
import norch.optim as optim
from norch.norchvision import transforms as T
import random
random.seed(1)
from memory_profiler import profile
@profile
def main():
BATCH_SIZE = 32
device = "cpu"
epochs = 10
transform = T.Compose(
[
T.ToTensor(),
T.Reshape([-1, 784, 1])
]
)
target_transform = T.Compose(
[
T.ToTensor()
]
)
print("Loading data")
train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)
train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE)
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(784, 30)
self.sigmoid1 = nn.Sigmoid()
self.fc2 = nn.Linear(30, 10)
self.sigmoid2 = nn.Sigmoid()
def forward(self, x):
out = self.fc1(x)
out = self.sigmoid1(out)
out = self.fc2(out)
out = self.sigmoid2(out)
return out
print("Creating model")
model = MyModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
loss_list = []
print("Starting training")
for epoch in range(epochs):
avg_loss = 0
num_steps = 0
for idx, batch in enumerate(train_loader):
if idx % 300 == 0 and idx > 0:
print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}")
break
inputs, target = batch
inputs = inputs.to(device)
target = target.to(device)
outputs = model(inputs)
loss = criterion(outputs, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss += loss[0]
num_steps += 1
break
avg_loss = avg_loss / num_steps
print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')
loss_list.append(avg_loss)
if __name__ == "__main__":
main()