{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Building Convolutional Neural Networks in JAX",
  "keywords": "",
  "domain": "图像处理",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/build-cnn-in-jax",
  "publish_date": "2025-03-12",
  "corpus_stage": "",
  "notes": "",
  "id": "19ace0a26b9c",
  "collected_at": "2026-10-02T14:11:15",
  "title": "Building Convolutional Neural Networks in JAX",
  "body_markdown": "Building Convolutional Neural Networks in JAX\nA comprehensive walkthrough of building a CNN using JAX, from data preprocessing to training and evaluation\nPublished\nMarch 12, 2025\n•\n9\nmin read\n•\nView as Markdown\nW\nWesley Kambale\nI'm a Machine Learning Engineer passionate about building production-ready ML systems for the African market.\nWith experience in TensorFlow, Keras, and Python-based workflows, I help teams bridge the gap between machine learning research and real-world deployment—especially on resource-constrained devices.\nI'm also a Google Developer Expert in AI. I regularly speak at tech conferences including PyCon Africa, DevFest Kampala, DevFest Nairobi and more and also write technical articles on AI/ML here.\nPart of series\nMachine Learning\nIntroduction\nDeep learning has revolutionized the field of artificial intelligence, and at the heart of this revolution are\nConvolutional Neural Networks (CNNs)\n. CNNs have become the go-to architectures for tasks involving images, such as object detection, facial recognition, medical imaging, and self-driving cars.\nTraditionally, frameworks like TensorFlow and PyTorch have dominated the deep learning landscape. However,\nJAX\nhas emerged as a powerful alternative, especially for research and high-performance computing. Developed by Google, JAX provides\nautomatic differentiation\nand\nJust-In-Time (JIT) compilation\n, making it highly efficient for numerical computing and deep learning applications.\nWhy Use JAX for CNNs?\nJAX stands out because of its ability to seamlessly run code on\nCPUs, GPUs, and TPUs\nwhile maintaining a NumPy-like API. This means you can develop models using familiar syntax while benefiting from:\nAutomatic Vectorization\n– With functions like\nvmap\n, JAX makes it easy to apply operations over large batches of data without writing explicit loops.\nEfficient Autograd\n– JAX provides\nautomatic differentiation\nusing\ngrad\n, which simplifies training deep learning models.\nXLA Compilation\n– Just-In-Time (JIT) compilation speeds up execution by compiling computation graphs for efficient hardware utilization.\nFunctional Programming Paradigm\n– Unlike traditional deep learning frameworks, JAX encourages\npure functions\n, which improves reproducibility and debugging.\nPrerequisites\nBefore proceeding, ensure you are familiar with:\nJAX fundamentals, you can check out the JAX documentation\nhere\n.\nBuilding CNNs in TensorFlow or PyTorch\nJAX optimizers and loss functions\nInstall Dependencies\nInstall the required libraries:\n!pip install jax jaxlib flax optax tensorflow tensorflow_datasets dm-pix tqdm matplotlib\njax\nand\njaxlib\n– The core JAX library and its hardware acceleration backend.\nflax\n– A neural network library for JAX, similar to PyTorch’s\ntorch.nn\n.\noptax\n– A library for optimization algorithms in JAX.\ndm_pix\n– A lightweight image processing library for JAX.\nmatplotlib\n– For visualizing images.\nImport Packages\nLoad necessary libraries:\nimport\nos\nimport\nnumpy\nas\nnp\nimport\ntensorflow\nas\ntf\nimport\ntensorflow_datasets\nas\ntfds\nimport\nmatplotlib.pyplot\nas\nplt\nimport\njax\nimport\njax.numpy\nas\njnp\nimport\noptax\nfrom\ntqdm.auto\nimport\ntqdm\nfrom\nflax\nimport\nlinen\nas\nnn\nfrom\nflax.training\nimport\ntrain_state\nimport\ndm_pix\nas\npix\n# Image processing in JAX\nVerify GPU Access\nJAX runs computations on\nCPUs, GPUs, and TPUs\nseamlessly. To check if your machine has a GPU or TPU available:\n# Get available devices\nprint(\n\"Available Devices:\"\n, jax.devices())\n# Check if GPU is available\nif\njax.default_backend() ==\n\"gpu\"\n:\n    print(\n\"Using GPU:\"\n, jax.devices(\n\"gpu\"\n))\nelif\njax.default_backend() ==\n\"tpu\"\n:\n    print(\n\"Using TPU:\"\n, jax.devices(\n\"tpu\"\n))\nelse\n:\n    print(\n\"Using CPU\"\n)\nData Preprocessing\nRaw image data comes in various sizes, orientations, and quality levels. Preprocessing is crucial for:\nEnsuring uniform input dimensions.\nNormalizing pixel values for stable training.\nAugmenting data to improve model generalization.\nConverting images into JAX-compatible tensors.\nLoading the Dataset\nWe will use the\nCats vs. Dogs dataset\navailable on Kaggle. Download and unzip the dataset using:\n!kaggle datasets download -d chetankv/dogs-cats-images\n!unzip dogs-cats-images.zip\nThis dataset gives us:\nA\ntraining dataset\n(\ntrain_data\n)\nA\ntest dataset\n(\ntest_data\n)\nDefine the path to the images and the batch size:\nbase_dir =\n\"/content/dog vs cat/dataset/training_set\"\nbatch_size =\n64\nResizing and Normalizing Images\nSince images in the dataset have varying sizes, we must\nresize\nthem to a fixed size (e.g.,\n128×128 pixels\n). Additionally, we normalize pixel values from\n[0, 255] → [0, 1]\nfor stable training.\nIMG_SIZE =\n128\nresize_and_rescale = tf.keras.Sequential(\n    [\n        tf.keras.layers.Resizing(IMG_SIZE, IMG_SIZE),\n        tf.keras.layers.Rescaling(\n1.0\n/\n255\n),\n    ]\n)\nData Augmentation for Better Generalization\nData augmentation helps improve model generalization by applying transformations like\nflipping, rotation, brightness adjustments, and cropping\n.\nrng = jax.random.PRNGKey(\n0\n)\nrng, inp_rng, init_rng = jax.random.split(rng,\n3\n)\n\ndelta =\n0.42\nfactor =\n0.42\n@jax.jit\ndef\ndata_augmentation\n(\nimage\n):\nnew_image = pix.adjust_brightness(image=image, delta=delta)\n    new_image = pix.random_brightness(image=new_image, max_delta=delta, key=inp_rng)\n    new_image = pix.flip_up_down(image=image)\n    new_image = pix.flip_left_right(image=new_image)\n    new_image = pix.rot90(k=\n1\n, image=new_image)\n# k = number of times the rotation is applied\nreturn\nnew_image\nConverting Data to JAX-Compatible Tensors\nJAX primarily operates on NumPy-like arrays (\njnp.array\n). TensorFlow uses\ntf.Tensor\n, so we must convert our dataset into a JAX-friendly format.\nAUTOTUNE = tf.data.AUTOTUNE\ndef\nprepare\n(\nds, shuffle=False\n):\n# Rescale and resize all datasets.\nds = ds.map(\nlambda\nx, y: (resize_and_rescale(x), y), num_parallel_calls=AUTOTUNE)\nif\nshuffle:\n        ds = ds.shuffle(\n1000\n)\n# Use buffered prefetching on all datasets.\nreturn\nds.prefetch(buffer_size=AUTOTUNE)\n\ntrain_ds = prepare(training_set, shuffle=\nTrue\n)\nval_ds = prepare(validation_set)\nevaluation_set = prepare(eval_set)\nVisualizing Preprocessed Images\nLet’s check if preprocessing is working as expected.\nplt.figure(figsize=(\n10\n,\n10\n))\n\naugmented_images = []\nfor\nimages, _\nin\ntraining_set.take(\n1\n):\nfor\ni\nin\nrange(\n9\n):\n    augmented_image = data_augmentation(np.array(images[i], dtype=jnp.float32))\n    augmented_images.append(augmented_image)\n    ax = plt.subplot(\n3\n,\n3\n, i +\n1\n)\n    plt.imshow(augmented_images[i].astype(\n\"uint8\"\n))\n    plt.axis(\n\"off\"\n)\nDefining a CNN in JAX\nIn JAX, neural networks are often implemented using\nFlax\n, a high-level neural network library designed to work seamlessly with JAX’s functional paradigm. Flax provides an intuitive way to define models using\nModule\nclasses.\nBelow is a simple implementation of a Convolutional Neural Network (CNN) in JAX using Flax:\nimport\njax.numpy\nas\njnp\nimport\nflax.linen\nas\nnn\nclass\nCNN\n(\nnn.Module\n):\nnum_classes: int\n# Number of output classes\n@nn.compact\ndef\n__call__\n(\nself, x\n):\nx = nn.Conv(features=\n32\n, kernel_size=(\n3\n,\n3\n), strides=(\n1\n,\n1\n), padding=\n\"SAME\"\n)(x)\n        x = nn.relu(x)\n        x = nn.max_pool(x, window_shape=(\n2\n,\n2\n), strides=(\n2\n,\n2\n))\n\n        x = nn.Conv(features=\n64\n, kernel_size=(\n3\n,\n3\n), strides=(\n1\n,\n1\n), padding=\n\"SAME\"\n)(x)\n        x = nn.relu(x)\n        x = nn.max_pool(x, window_shape=(\n2\n,\n2\n), strides=(\n2\n,\n2\n))\n\n        x = nn.Conv(features=\n128\n, kernel_size=(\n3\n,\n3\n), strides=(\n1\n,\n1\n), padding=\n\"SAME\"\n)(x)\n        x = nn.relu(x)\n        x = nn.max_pool(x, window_shape=(\n2\n,\n2\n), strides=(\n2\n,\n2\n))\n\n        x = x.reshape((x.shape[\n0\n],\n-1\n))\n# Flatten feature maps\nx = nn.Dense(features=\n128\n)(x)\n        x = nn.relu(x)\n\n        x = nn.Dense(features=self.num_classes)(x)\n# Output layer\nreturn\nx\nFirst Convolutional Layer\nApplies a\n32-channel\nconvolution with a\n3×3\nkernel.\nUses\nReLU\nactivation to introduce non-linearity.\nApplies\nmax pooling\nwith a\n2×2\nwindow and a stride of\n2\n, reducing the spatial dimensions by half.\nSecond Convolutional Layer\nUses a\n64-channel\nconvolution with a\n3×3\nkernel.\nAgain applies\nReLU\nactivation.\nAnother\nmax pooling\noperation further reduces spatial dimensions.\nThird Convolutional Layer\nIncreases the number of channels to\n128\nwhile keeping the\n3×3\nkernel size.\nApplies\nReLU\nactivation and another\nmax pooling\nstep.\nFlattening and Fully Connected Layers\nThe feature maps from the final convolutional layer are\nflattened\ninto a 1D vector.\nA\ndense layer with 128 neurons\napplies a ReLU activation.\nThe final\noutput layer\nproduces logits corresponding to the number of classes.\nWhy Use\n@nn.compact\n?\nFlax provides two ways to define models:\nUsing\n@nn.compact\n, which allows direct instantiation of layers within the\n__call__\nmethod.\nUsing\nsetup()\n, where layers are explicitly defined as attributes.\nThe\ncompact\napproach is cleaner and more intuitive for simple models, avoiding the need to define layer attributes separately.\nInitializing the Model\nJAX does not use an implicit state, so model parameters must be explicitly initialized. The\ninit\nfunction from Flax helps generate the model’s parameters using a random key and an input shape.\nfrom\nflax.core\nimport\nfreeze, unfreeze\n# Set up PRNG key\nkey = jax.random.PRNGKey(\n0\n)\n# Define input shape (batch_size, height, width, channels)\ninput_shape = (\n1\n,\n32\n,\n32\n,\n3\n)\n# Example for a 32x32 RGB image\n# Initialize model\nmodel = CNN(num_classes=\n10\n)\n# Assuming 10 output classes\nparams = model.init(key, jnp.ones(input_shape))[\n\"params\"\n]\nA\nrandom key\nis generated using\njax.random.PRNGKey(0)\n. JAX requires explicit control over random number generation for reproducibility.\nA\ndummy input tensor\nof shape\n(1, 32, 32, 3)\nis created to initialize the network.\nThe model is\ninstantiated\nand the\ninit\nfunction generates model parameters using the random key.\nThe\n\"params\"\nfield is extracted from the initialization output, as Flax’s\ninit\nmethod returns a dictionary containing additional information (e.g., batch statistics if using BatchNorm).\nDefining the Training State\nFlax provides a\ntrain_state\nabstraction to manage model parameters, optimizer state, and other training-related information. The\noptax\nlibrary is used for defining the optimizer.\nimport\noptax\nfrom\nflax.training\nimport\ntrain_state\nclass\nTrainState\n(\ntrain_state.TrainState\n):\npass\n# No additional attributes needed for now\n# Define the optimizer\nlearning_rate =\n0.001\noptimizer = optax.adam(learning_rate)\n# Initialize the training state\nstate = TrainState.create(apply_fn=model.apply, params=params, tx=optimizer)\nTrainState\nis a dataclass that stores the model's parameters, optimizer state, and\napply_fn\n(the function used for forward passes).\nOptax's Adam optimizer\nis set up with a learning rate of\n0.001\n.\nThe\nstate.create()\nmethod initializes the model’s training state with:\napply_fn\n: The forward pass function from the model.\nparams\n: The initialized parameters from the previous step.\ntx\n: The optimizer (Adam in this case).\nDefining Loss and Accuracy Metrics\nA loss function is required to guide training, while an accuracy function evaluates model performance.\nLoss Function\nimport\njax.nn\nas\njnn\ndef\ncross_entropy_loss\n(\nparams, state, batch\n):\nlogits = state.apply_fn({\n'params'\n: params}, batch[\n'images'\n])\n    labels = jnn.one_hot(batch[\n'labels'\n], num_classes=\n10\n)\nreturn\n-jnp.sum(labels * jnn.log_softmax(logits)) / batch[\n'labels'\n].shape[\n0\n]\nThe function takes\nmodel parameters\n, the\ncurrent training state\n, and a\nbatch of input data\n.\nThe\nforward pass\nis performed using\napply_fn\n, producing logits (raw model predictions).\nThe\nlabels are one-hot encoded\nto match the logits' shape.\nThe\ncross-entropy loss\nis computed using\nlog_softmax(logits)\n, ensuring numerical stability.\nThe loss is\naveraged over the batch size\nfor proper optimization.\nAccuracy Function\ndef\ncompute_accuracy\n(\nparams, state, batch\n):\nlogits = state.apply_fn({\n'params'\n: params}, batch[\n'images'\n])\n    predicted_labels = jnp.argmax(logits, axis=\n1\n)\nreturn\njnp.mean(predicted_labels == batch[\n'labels'\n])\nThe function takes\nmodel parameters\n, the\ntraining state\n, and a\nbatch of data\n.\nThe\nforward pass\nis executed to obtain logits.\nThe\nhighest-scoring class\nis selected using\nargmax()\n, determining the model’s predicted label.\nAccuracy is computed by\ncomparing predictions with actual labels\nand averaging the correct classifications.\nTraining and Evaluating a CNN in JAX\nTraining in JAX is based on functional transformations, meaning explicit gradient computation and parameter updates are required. The\njax.grad\nfunction is used to compute gradients efficiently.\nTraining Step Function\n@jax.jit\ndef\ntrain_step\n(\nstate, batch\n):\ndef\nloss_fn\n(\nparams\n):\nlogits = state.apply_fn({\n'params'\n: params}, batch[\n'images'\n])\n        labels = jax.nn.one_hot(batch[\n'labels'\n], num_classes=\n10\n)\n        loss = -jnp.sum(labels * jax.nn.log_softmax(logits)) / batch[\n'labels'\n].shape[\n0\n]\nreturn\nloss\n# Compute gradients\ngrads = jax.grad(loss_fn)(state.params)\n# Update model state\nstate = state.apply_gradients(grads=grads)\nreturn\nstate\nJIT Compilation (\n@jax.jit\n): JAX’s just-in-time compilation speeds up training by optimizing computation.\nloss_fn\nFunction: Defines the cross-entropy loss to be minimized.\njax.grad(loss_fn)\n: Computes gradients with respect to model parameters.\nstate.apply_gradients(grads=grads)\n: Updates the training state using computed gradients.\nTraining Loop\ndef\ntrain_model\n(\nstate, train_loader, num_epochs=\n10\n):\nfor\nepoch\nin\nrange(num_epochs):\nfor\nbatch\nin\ntrain_loader:\n            state = train_step(state, batch)\n        print(\nf\"Epoch\n{epoch +\n1\n}\ncompleted\"\n)\nreturn\nstate\nIterates through multiple epochs\n, training the model for\nnum_epochs\n.\nProcesses each batch\n, updating the model parameters.\nLogs progress\nat the end of each epoch.\nEvaluation Step Function\ndef\nevaluate_model\n(\nstate, test_loader\n):\naccuracies = []\nfor\nbatch\nin\ntest_loader:\n        acc = compute_accuracy(state.params, state, batch)\n        accuracies.append(acc)\n\n    final_accuracy = jnp.mean(jnp.array(accuracies))\n    print(\nf\"Test Accuracy:\n{final_accuracy *\n100\n:\n.2\nf}\n%\"\n)\nreturn\nfinal_accuracy\nIterates through the\ntest dataset\n, computing accuracy for each batch.\nAggregates accuracy scores\nacross batches to compute the final accuracy.\nPrints the test accuracy\n, indicating how well the model generalizes to unseen data.\nConclusion\nJAX’s functional and hardware-accelerated approach allows for efficient model training, particularly on GPUs and TPUs. The explicit handling of gradients and optimizers ensures flexibility while maintaining high performance.\nFuture work could explore\nadvanced techniques\nsuch as data augmentation, model regularization, and hyperparameter tuning to improve performance. Additionally, integrating JAX with frameworks like TensorFlow or PyTorch could provide hybrid workflows for deep learning research and production applications.\n#\ncnn\n#\njax\n#\ncnn-for-begginers",
  "metadata": {
    "title": "Building Convolutional Neural Networks in JAX",
    "description": "A comprehensive walkthrough of building a CNN using JAX, from data preprocessing to training and evaluation",
    "keywords": null,
    "author": null,
    "og:title": "Building Convolutional Neural Networks in JAX",
    "og:description": "A comprehensive walkthrough of building a CNN using JAX, from data preprocessing to training and evaluation",
    "og:url": "https://kambale.dev/build-cnn-in-jax",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1741791803057/7b61f8c9-6ca5-4564-9bf4-5b385efaf4dc.png",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:creator": "https://twitter.com/WesKambale",
    "twitter:title": "Building Convolutional Neural Networks in JAX",
    "twitter:description": "A comprehensive walkthrough of building a CNN using JAX, from data preprocessing to training and evaluation",
    "twitter:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1741791803057/7b61f8c9-6ca5-4564-9bf4-5b385efaf4dc.png",
    "article:published_time": "2025-03-12T15:05:24.838Z",
    "article:author": "https://kambale.dev/@kambale"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/19ace0a26b9c/code/body-000.txt",
      "chars": 89,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-001.txt",
      "chars": 310,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-002.txt",
      "chars": 285,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-003.txt",
      "chars": 83,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-004.txt",
      "chars": 70,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-005.txt",
      "chars": 172,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-006.txt",
      "chars": 510,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-007.txt",
      "chars": 443,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-008.txt",
      "chars": 344,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-009.txt",
      "chars": 903,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-010.txt",
      "chars": 344,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-011.txt",
      "chars": 340,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-012.txt",
      "chars": 272,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-013.txt",
      "chars": 216,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-014.txt",
      "chars": 466,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-015.txt",
      "chars": 232,
      "language": "text"
    },
    {
      "file": "attachments/19ace0a26b9c/code/body-016.txt",
      "chars": 318,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/res/hashnode/image/upload/v1741790696877/2599696c-54c6-4121-9022-c158009ae0bd.png",
      "bytes": 116831,
      "reused": true,
      "local_path": "attachments/19ace0a26b9c/images/000.png",
      "src": "https://cdn.hashnode.com/res/hashnode/image/upload/v1741790696877/2599696c-54c6-4121-9022-c158009ae0bd.png",
      "alt": "Building Convolutional Neural Networks in JAX",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/build-cnn-in-jax.md",
      "label": "View as Markdown",
      "type": "a"
    },
    {
      "href": "https://hashnode.com/@kambale",
      "label": "W",
      "type": "a"
    },
    {
      "href": "https://twitter.com/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://linkedin.com/in/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/series/ml-series",
      "label": "Part of series Machine Learning",
      "type": "a"
    },
    {
      "href": "https://docs.jax.dev/en/latest/index.html",
      "label": "here",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/cnn",
      "label": "# cnn",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/cnn-for-begginers",
      "label": "# cnn-for-begginers",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:48Z",
  "source_url": "https://kambale.dev/build-cnn-in-jax",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 14959,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/19ace0a26b9c/pages/source.html",
  "local_html": "attachments/19ace0a26b9c/pages/article.html",
  "body_base_path": "attachments/19ace0a26b9c/pages/",
  "body_markdown_path": "attachments/19ace0a26b9c/pages/article.md"
}