Skip to content

Commit

Permalink
“referece_tutorial_links_added”
Browse files Browse the repository at this point in the history
  • Loading branch information
@selamw1
selamw1 committed Dec 3, 2024
1 parent f6b8e39 commit f97b3d9
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 49 deletions.
88 changes: 46 additions & 42 deletions docs/data_loaders_on_gpu_with_jax.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,12 @@
"\n",
"You'll see how to use each of these libraries to efficiently load data for a simple image classification task using the MNIST dataset.\n",
"\n",
"Compared to the [Data Loaders on CPU](https://jax-ai-stack.readthedocs.io/en/latest/data_loaders_on_cpu_with_jax.html), working with GPUs introduces opportunities for further optimization, such as transferring data to the GPU using `device_put`, leveraging larger batch sizes for faster processing, and addressing considerations like memory management."
"Compared to [CPU-based loading](https://jax-ai-stack.readthedocs.io/en/latest/data_loaders_on_cpu_with_jax.html), working with a GPU introduces specific considerations like transferring data to the GPU using `device_put`, managing larger batch sizes for faster processing, and efficiently utilizing GPU memory. Unlike multi-device setups, this guide focuses on optimizing data handling for a single GPU.\n",
"\n",
"\n",
"If you're looking for CPU-specific data loading advice, see [Data Loaders on CPU](https://jax-ai-stack.readthedocs.io/en/latest/data_loaders_on_cpu_with_jax.html).\n",
"\n",
"If you're looking for a multi-device data loading strategy, see [Data Loaders on Multi-Device Setups](https://jax-ai-stack.readthedocs.io/en/latest/data_loaders_for_multi_device_setups_with_jax.html)."
]
},
{
Expand All @@ -34,12 +39,12 @@
"id": "-rsMgVtO6asW"
},
"source": [
"### Import JAX API"
"## Import JAX API"
]
},
{
"cell_type": "code",
"execution_count": 35,
"execution_count": null,
"metadata": {
"id": "tDJNQ6V-Dg5g"
},
Expand All @@ -56,12 +61,12 @@
"id": "TsFdlkSZKp9S"
},
"source": [
"### Checking GPU Availability for JAX"
"## Checking GPU Availability for JAX"
]
},
{
"cell_type": "code",
"execution_count": 36,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -91,14 +96,14 @@
"id": "qyJ_WTghDnIc"
},
"source": [
"### Setting Hyperparameters and Initializing Parameters\n",
"## Setting Hyperparameters and Initializing Parameters\n",
"\n",
"You'll define hyperparameters for your model and data loading, including layer sizes, learning rate, batch size, and the data directory. You'll also initialize the weights and biases for a fully-connected neural network."
]
},
{
"cell_type": "code",
"execution_count": 37,
"execution_count": null,
"metadata": {
"id": "qLNOSloFDka_"
},
Expand Down Expand Up @@ -133,7 +138,7 @@
"id": "rHLdqeI7D2WZ"
},
"source": [
"### Model Prediction with Auto-Batching\n",
"## Model Prediction with Auto-Batching\n",
"\n",
"In this section, you'll define the `predict` function for your neural network. This function computes the output of the network for a single input image.\n",
"\n",
Expand All @@ -142,7 +147,7 @@
},
{
"cell_type": "code",
"execution_count": 38,
"execution_count": null,
"metadata": {
"id": "bKIYPSkvD1QV"
},
Expand Down Expand Up @@ -174,7 +179,7 @@
"id": "rLqfeORsERek"
},
"source": [
"### Utility and Loss Functions\n",
"## Utility and Loss Functions\n",
"\n",
"You'll now define utility functions for:\n",
"- One-hot encoding: Converts class indices to binary vectors.\n",
Expand All @@ -190,7 +195,7 @@
},
{
"cell_type": "code",
"execution_count": 39,
"execution_count": null,
"metadata": {
"id": "sA0a06raEQfS"
},
Expand Down Expand Up @@ -253,7 +258,7 @@
},
{
"cell_type": "code",
"execution_count": 40,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -287,7 +292,7 @@
},
{
"cell_type": "code",
"execution_count": 41,
"execution_count": null,
"metadata": {
"id": "kO5_WzwY59gE"
},
Expand All @@ -301,7 +306,7 @@
},
{
"cell_type": "code",
"execution_count": 42,
"execution_count": null,
"metadata": {
"id": "6f6qU8PCc143"
},
Expand Down Expand Up @@ -348,7 +353,7 @@
},
{
"cell_type": "code",
"execution_count": 43,
"execution_count": null,
"metadata": {
"id": "Kxbl6bcx6crv"
},
Expand All @@ -370,7 +375,7 @@
},
{
"cell_type": "code",
"execution_count": 44,
"execution_count": null,
"metadata": {
"id": "c9ZCJq_rzPck"
},
Expand All @@ -393,7 +398,7 @@
},
{
"cell_type": "code",
"execution_count": 45,
"execution_count": null,
"metadata": {
"id": "brlLG4SqGphm"
},
Expand All @@ -406,7 +411,7 @@
},
{
"cell_type": "code",
"execution_count": 46,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -446,7 +451,7 @@
},
{
"cell_type": "code",
"execution_count": 47,
"execution_count": null,
"metadata": {
"id": "0LdT8P8aisWF"
},
Expand All @@ -469,7 +474,7 @@
},
{
"cell_type": "code",
"execution_count": 48,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -510,7 +515,7 @@
},
{
"cell_type": "code",
"execution_count": 49,
"execution_count": null,
"metadata": {
"id": "sGaQAk1DHMUx"
},
Expand All @@ -532,7 +537,7 @@
},
{
"cell_type": "code",
"execution_count": 50,
"execution_count": null,
"metadata": {
"id": "1hOamw_7C8Pb"
},
Expand All @@ -556,7 +561,7 @@
},
{
"cell_type": "code",
"execution_count": 51,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -592,7 +597,7 @@
},
{
"cell_type": "code",
"execution_count": 52,
"execution_count": null,
"metadata": {
"id": "vX59u8CqEf4J"
},
Expand Down Expand Up @@ -620,7 +625,7 @@
},
{
"cell_type": "code",
"execution_count": 53,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -670,7 +675,7 @@
},
{
"cell_type": "code",
"execution_count": 54,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -714,7 +719,7 @@
},
{
"cell_type": "code",
"execution_count": 55,
"execution_count": null,
"metadata": {
"id": "mS62eVL9Ifmz"
},
Expand All @@ -738,7 +743,7 @@
},
{
"cell_type": "code",
"execution_count": 56,
"execution_count": null,
"metadata": {
"id": "bnrhac5Hh7y1"
},
Expand Down Expand Up @@ -772,7 +777,7 @@
},
{
"cell_type": "code",
"execution_count": 57,
"execution_count": null,
"metadata": {
"id": "pN3oF7-ostGE"
},
Expand All @@ -792,7 +797,7 @@
},
{
"cell_type": "code",
"execution_count": 58,
"execution_count": null,
"metadata": {
"id": "f1VnTuX3u_kL"
},
Expand All @@ -810,7 +815,7 @@
},
{
"cell_type": "code",
"execution_count": 59,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -844,7 +849,7 @@
},
{
"cell_type": "code",
"execution_count": 60,
"execution_count": null,
"metadata": {
"id": "2jqd1jJt25Bj"
},
Expand Down Expand Up @@ -875,7 +880,7 @@
},
{
"cell_type": "code",
"execution_count": 61,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -925,7 +930,7 @@
},
{
"cell_type": "code",
"execution_count": 62,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -979,7 +984,7 @@
},
{
"cell_type": "code",
"execution_count": 63,
"execution_count": null,
"metadata": {
"id": "8v1N59p76zn0"
},
Expand All @@ -999,7 +1004,7 @@
},
{
"cell_type": "code",
"execution_count": 64,
"execution_count": null,
"metadata": {
"id": "a22kTvgk6_fJ"
},
Expand All @@ -1021,7 +1026,7 @@
},
{
"cell_type": "code",
"execution_count": 65,
"execution_count": null,
"metadata": {
"id": "NHrKatD_7HbH"
},
Expand All @@ -1047,7 +1052,7 @@
},
{
"cell_type": "code",
"execution_count": 66,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -1083,7 +1088,7 @@
},
{
"cell_type": "code",
"execution_count": 67,
"execution_count": null,
"metadata": {
"id": "-zLJhogj7RL-"
},
Expand All @@ -1109,7 +1114,7 @@
},
{
"cell_type": "code",
"execution_count": 68,
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
Expand Down Expand Up @@ -1145,15 +1150,14 @@
"source": [
"## Summary\n",
"\n",
"This notebook explored efficient methods for loading data on a GPU with JAX, using libraries like PyTorch DataLoader, TensorFlow Datasets, Grain, and Hugging Face Datasets. You also learned GPU-specific optimizations, such as `device_put` for data transfer and memory management, to enhance training efficiency. Each methods offers unique benefits, helping you choose the best fit for your project needs."
"This notebook explored efficient methods for loading data on a GPU with JAX, using libraries such as PyTorch DataLoader, TensorFlow Datasets, Grain, and Hugging Face Datasets. You also learned GPU-specific optimizations, including using `device_put` for data transfer and managing GPU memory, to enhance training efficiency. Each method offers unique benefits, allowing you to choose the best approach based on your project requirements."
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"name": "data_loaders_on_gpu_with_jax.ipynb",
"provenance": []
},
"jupytext": {
Expand Down
Loading

0 comments on commit f97b3d9

Please sign in to comment.