diff --git a/build/cuda.cu.o b/build/cuda.cu.o index 17d7461..437afeb 100644 Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ diff --git a/build/tensor.o b/build/tensor.o index 09cfc55..7b38f19 100644 Binary files a/build/tensor.o and b/build/tensor.o differ diff --git a/examples/train_multigpu.py b/examples/train_multigpu.py index 8cbaaff..8d112cb 100644 --- a/examples/train_multigpu.py +++ b/examples/train_multigpu.py @@ -38,6 +38,7 @@ def main(): ] ) + print("Loading data") train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform) distributed_sampler = DistributedSampler(dataset=train_data, num_replicas=world_size, rank=rank) train_loader = norch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, sampler=distributed_sampler) @@ -58,6 +59,7 @@ def main(): return out + print("Creating model") model = MyModel().to(device) model = DistributedDataParallel(model) criterion = nn.CrossEntropyLoss() @@ -67,8 +69,15 @@ def main(): print(f"Starting training on Rank {rank}/{world_size}\n\n") for epoch in range(epochs): + + avg_loss = 0 + num_steps = 0 + for idx, batch in enumerate(train_loader): + if idx % 100 == 0 and rank == 0: + print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}") + inputs, target = batch inputs = inputs.to(device) @@ -83,10 +92,14 @@ def main(): loss.backward() optimizer.step() - + + avg_loss += loss[0] + num_steps += 1 + + avg_loss = avg_loss / num_steps if rank == 0: - print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}') - loss_list.append(loss[0]) + print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}') + loss_list.append(avg_loss) if __name__ == "__main__": main() diff --git a/examples/train_nn_mnist.ipynb b/examples/train_nn_mnist.ipynb index 21e61f7..12c37ce 100644 --- a/examples/train_nn_mnist.ipynb +++ b/examples/train_nn_mnist.ipynb @@ -139,7 +139,11 @@ "optimizer = optim.SGD(model.parameters(), lr=0.01)\n", "loss_list = []\n", "\n", - "for epoch in range(epochs): \n", + "for epoch in range(epochs): \n", + " \n", + " avg_loss = 0 \n", + " num_steps = 0\n", + " \n", " for idx, batch in enumerate(train_loader):\n", "\n", " inputs, target = batch\n", @@ -157,8 +161,12 @@ "\n", " optimizer.step()\n", "\n", - " print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss[0]:.4f}')\n", - " loss_list.append(loss[0])" + " avg_loss += loss[0]\n", + " num_steps += 1\n", + "\n", + " avg_loss = avg_loss / num_steps\n", + " print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}')\n", + " loss_list.append(avg_loss)" ] }, { diff --git a/examples/train_singlegpu.py b/examples/train_singlegpu.py index 8c58ed6..6b0dfac 100644 --- a/examples/train_singlegpu.py +++ b/examples/train_singlegpu.py @@ -24,6 +24,8 @@ def main(): ] ) + + 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) @@ -43,14 +45,23 @@ def main(): return out + print("Creating model") model = MyModel().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.01) loss_list = [] - for epoch in range(epochs): + print("Starting training") + for epoch in range(epochs): + + avg_loss = 0 + num_steps = 0 + for idx, batch in enumerate(train_loader): + if idx % 100 == 0: + print(f"Epoch: {epoch}/{epochs} - Step: {idx} / {len(train_loader)}") + inputs, target = batch inputs = inputs.to(device) @@ -64,6 +75,13 @@ def main(): loss.backward() optimizer.step() + + avg_loss += loss[0] + num_steps += 1 + + avg_loss = avg_loss / num_steps + print(f'Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}') + loss_list.append(avg_loss) if __name__ == "__main__": diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc index 3662c85..ad61d69 100644 Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp index e188e7e..54882bd 100644 --- a/norch/csrc/tensor.cpp +++ b/norch/csrc/tensor.cpp @@ -26,23 +26,29 @@ extern "C" { tensor->device = strdup(device); tensor->strides = (int*)malloc(ndim * sizeof(int)); tensor->shape = (int*)malloc(ndim * sizeof(int)); - tensor->data = (float*)malloc(tensor->size * sizeof(float)); - if (tensor->device == NULL || tensor->strides == NULL || tensor->shape == NULL || tensor->data == NULL) { + if (tensor->device == NULL || tensor->strides == NULL || tensor->shape == NULL) { fprintf(stderr, "Memory allocation failed\n"); exit(1); } memcpy(tensor->shape, shape, tensor->ndim * sizeof(int)); strcpy(tensor->device, device); - memcpy(tensor->data, data, tensor->size * sizeof(float)); - + + if (strcmp(tensor->device, "cpu") == 0) { + tensor->data = (float*)malloc(tensor->size * sizeof(float)); + memcpy(tensor->data, data, tensor->size * sizeof(float)); + } else { + //already allocated on gpu + tensor->data = data; + } + int stride = 1; for (int i = ndim - 1; i >= 0; i--) { tensor->strides[i] = stride; stride *= shape[i]; } - + return tensor; } diff --git a/norch/libtensor.so b/norch/libtensor.so index d7435de..300caf1 100755 Binary files a/norch/libtensor.so and b/norch/libtensor.so differ diff --git a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc index 0daa1ff..6897844 100644 Binary files a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc and b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc differ diff --git a/norch/tensor.py b/norch/tensor.py index 22061be..0287fb6 100644 --- a/norch/tensor.py +++ b/norch/tensor.py @@ -46,13 +46,16 @@ class Tensor: Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_char_p] Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor) - self.tensor = Tensor._C.create_tensor( self.data_ctype, self.shape_ctype, self.ndim_ctype, self.device_ctype ) + + del self.data_ctype + del self.shape_ctype + del self.device_ctype else: self.tensor = None,