How To Build a Production-Grade Deep Learning Pipeline with Weights & Biases

A near-to-production example using amino acid sequence modeling

Photo by Steve Johnson on Unsplash

Most deep learning projects fail not because the model architecture was wrong, but because the surrounding infrastructure was too fragile to iterate on. Data pipelines break silently. Hyperparameter searches happen in ad hoc notebook cells that cannot be reproduced. Training runs on a GPU cluster diverge from local development because the environments were never truly equivalent. Results live in a researcher’s local TensorBoard instance and can’t be shared without exporting CSVs.

These are not exotic problems. They are the default state of a project that grows organically without deliberate MLOps thinking from the start. The cost compounds: debugging a silent data schema violation three weeks after it was introduced is an order of magnitude more expensive than catching it at ingestion time. A hyperparameter sweep that cannot be reproduced is scientific debt.

Here, I am going to walk you though one of my side projects,

[embed]GitHub - Dima806/antibody_sequence_modeling_example: Example project for predicting biophysical… Example project for predicting biophysical properties of antibody CDR (Complementarity-Determining Region) sequences …github.com

that can work as a reference implementation of what solution looks like.

One important caveat, though: this is not a fully production-ready system. It imitates production-grade patterns, but it retains one intentionally manual step: GPU training is triggered by hand inside a Kaggle Notebook,

[embed]train_OAS_data_WANDB_example Explore and run machine learning code with Kaggle Notebooks | Using data from No attached data sourceswww.kaggle.com

As a result, there is no CI/CD pipeline that automatically dispatches training jobs to a GPU cluster. However, the patterns it does implement (validated data contracts, Makefile-driven operations, reproducible sweep configurations, and continuous experiment visibility through Weights & Biases) are the same you would find in a fully automated system. This project teaches you how to think about them, and it gives you working code you can run today. Replacing the manual Kaggle step with a cloud training dispatch (Amazon’s SageMaker, Google’s Vertex AI, Microsoft’s Azure ML, or similar) would be a targeted addition, not a re-architecture. But that step is left explicit and manual here, because the audience for this project is individual researchers and small teams who do not have a GPU machine sitting on their desk and perhaps do not want to pay for costly cloud compute while iterating on code.

Having said that, let’s dive in!

The OAS dataset: a credible real-world problem

A pipeline reference implementation needs a real problem to be credible. Toy datasets obscure the engineering challenges that arise at scale — large file downloads, schema heterogeneity, the need for feature extraction, class imbalance, and the tradeoffs between smoke-scale and production-scale training. Antibody CDR-H3 sequence modeling, using sampled data from the Observed Antibody Space (OAS):

[embed]OAS: Observed Antibody Space The Observed Antibody Space database, or OAS, is a project to collect and annotate immune repertoires for use in…opig.stats.ox.ac.uk

provides all of these challenges in a domain that is scientifically meaningful.

This is, in other words, a real ML problem with real data engineering challenges — which is exactly what makes it a useful vehicle for demonstrating production-grade pipeline design. Let us walk through how it is being built, how it runs, and why each design decision matters.

Repository architecture: two environments, one manual bridge

One of the most deliberate design choices in this project is the dual-environment strategy: the codebase is designed to run identically in both a lightweight CPU environment and a full GPU environment.

Development happens in GitHub Codespaces. Codespaces provides a free, reproducible cloud development environment, but the free tier is limited to 2 CPUs and no GPU. That is perfectly adequate for writing code, running tests, linting, generating synthetic smoke data, and doing short smoke training runs. It is not adequate for training on 500K sequences with a deep Transformer or ESM-2. This is where the vast majority of iteration happens — fast, cheap, and reproducible.

Full training happens in a Kaggle Notebook. Kaggle provides free GPU access (T4 or P100) with no cloud billing, but the interface is a Jupyter notebook that you open in a browser and run manually. There is no API to dispatch a training job programmatically from your development machine. This is the intentional manual seam in the pipeline: you open the notebook, run the cells, and the results flow automatically into W&B. It is not CI/CD. It is a deliberate, human-in-the-loop step.

This is the tradeoff of a zero-cost GPU training setup. A fully automated system would replace the Kaggle notebook with a cloud training (SageMaker, Vertex AI, or similar) triggered automatically on a git push or via a Makefile target. That would eliminate the manual step entirely. This project does not do that, and it is worth saying so plainly. What it does instead is make the manual step as clean and reproducible as possible: a notebook that self-bootstraps, runs the same make commands as your development machine, and pushes results to the same W&B project. The gap between this and full automation is a single integration point, not a different philosophy.

Understanding this constraint also clarifies the purpose of the smoke workflow. Because the Kaggle GPU run is a manual, relatively costly step (even at zero monetary cost, it takes time), the project makes local CPU iteration fast and reliable. The smoke dataset, smoke training targets, and CPU sweep exist specifically so that you can validate the entire pipeline end-to-end on a 2-CPU GitHub Codespaces machine in several minutes, before committing to a full-scale GPU run. This is exactly the kind of discipline that a production system enforces at scale, even if here it is enforced by convention rather than by infrastructure:

antibody_sequence_modeling_example/
├── configs/
│   ├── default.yaml          # Smoke-scale defaults
│   ├── schema.yaml           # Dataset schema (validated on every run)
│   ├── sweep_cpu.yaml        # Random sweep, 5 runs, CPU
│   └── sweep_gpu.yaml        # Bayesian sweep, 50 runs, GPU
├── data/
│   ├── smoke/
│   │   └── sequences_smoke.csv   # ~2K synthetic sequences (committed)
│   ├── download.py           # Download OAS data from Zenodo / OAS API
│   ├── generate_smoke.py     # Generate synthetic smoke dataset
│   └── validate.py           # Schema validation
├── src/antibody_seq_ml/
│   ├── dataset.py            # Tokeniser, CDRDataset, DataLoaders
│   ├── train.py              # Training loop + W&B logging
│   ├── evaluate.py           # Metrics, confusion matrix, attention viz
│   ├── sweep.py              # W&B sweep agent entry point
│   └── models/
│       ├── base.py           # SequenceModel base class + factory
│       ├── lstm.py           # BiLSTM with mean pooling
│       ├── transformer.py    # TransformerEncoder + sinusoidal PE
│       └── esm2.py           # ESM-2 fine-tuning (GPU only)
├── tests/
├── Makefile
└── pyproject.toml

The --smoke flag threads through the entire stack, switching between small-scale synthetic data (~2K sequences, fast CPU epochs) and larger-scale OAS sample data (~500K sequences, GPU-optimised batching). The same Python code runs in both cases. Hardware is never hardcoded: every model and training script branches on torch.cuda.is_available().

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model_factory(cfg).to(device)

This single pattern, applied consistently, means there are no divergent code paths to maintain and no surprises when the environment changes.

The dataset and schema validation

Training data flows through a validated pipeline. After every make data-smoke or make data-full call, the output CSV is automatically checked against configs/schema.yaml before anything downstream consumes it. This catches bad downloads, silent API changes, and data drift before they silently corrupt model training.

The schema enforces six required columns, value ranges, string length bounds, and cross-column consistency:

# configs/schema.yaml
columns:
  sequence_id:    {dtype: str,   nullable: false}
  heavy_chain_sequence: {dtype: str, min_length: 30, nullable: false}
  cdr_h3:         {dtype: str,   min_length: 5, max_length: 30, nullable: false}
  cdr_h3_length:  {dtype: int,   nullable: false}
  length_class:   {dtype: str,   values: [short, medium, long], nullable: false}
  hydrophobicity: {dtype: float, min: -5.0, max: 5.0, nullable: false}

cross_checks:
  - cdr_h3_length == len(cdr_h3)
  - length_class == "short"  iff cdr_h3_length <= 9
  - length_class == "medium" iff 10 <= cdr_h3_length <= 14
  - length_class == "long"   iff cdr_h3_length >= 15

This kind of schema enforcement at the data boundary is one of those practices that feels like overhead until the day it catches a silent bug that would otherwise have taken hours to diagnose. In a research-to-production pipeline, it’s non-negotiable.

Three model architectures

The project compares three architectures of increasing complexity and capability, all built on a shared SequenceModel base class with a factory function for clean instantiation:

def model_factory(cfg: DictConfig) -> SequenceModel:
    model_type = cfg.model.type
    if model_type == "lstm":
        return BiLSTMModel(cfg)
    elif model_type == "transformer":
        return TransformerModel(cfg)
    elif model_type == "esm2":
        return ESM2Model(cfg)
    else:
        raise ValueError(f"Unknown model type: {model_type}")

BiLSTM is the fast, interpretable baseline. Amino acid sequences are embedded, processed bidirectionally through an LSTM, and mean-pooled into a fixed-length representation:

class BiLSTMModel(SequenceModel):
    def __init__(self, cfg: DictConfig) -> None:
        super().__init__()
        self.embedding = nn.Embedding(cfg.vocab_size, cfg.model.embedding_dim)
        self.lstm = nn.LSTM(
            cfg.model.embedding_dim,
            cfg.model.hidden_dim,
            num_layers=cfg.model.num_layers,
            bidirectional=True,
            batch_first=True,
            dropout=cfg.model.dropout,
        )
        feat_dim = cfg.model.hidden_dim * 2
        self.cls_head = nn.Linear(feat_dim, 3)    # short / medium / long
        self.reg_head = nn.Linear(feat_dim, 1)    # hydrophobicity GRAVY

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        emb = self.embedding(x)
        out, _ = self.lstm(emb)
        pooled = out.mean(dim=1)
        return self.cls_head(pooled), self.reg_head(pooled).squeeze(-1)

Transformer replaces the recurrence with self-attention. Here, a[CLS] token is prepended to the sequence, sinusoidal positional encodings are added, and the CLS output is passed to both task heads:

class TransformerModel(SequenceModel):
    def __init__(self, cfg: DictConfig) -> None:
        super().__init__()
        d = cfg.model.d_model
        self.embedding  = nn.Embedding(cfg.vocab_size + 1, d)  # +1 for CLS
        self.pos_enc    = SinusoidalPositionalEncoding(d, cfg.model.max_seq_length)
        encoder_layer   = nn.TransformerEncoderLayer(
            d_model=d, nhead=cfg.model.nhead,
            dim_feedforward=cfg.model.dim_feedforward,
            dropout=cfg.model.dropout, batch_first=True,
        )
        self.encoder    = nn.TransformerEncoder(encoder_layer, cfg.model.num_layers)
        self.cls_head   = nn.Linear(d, 3)
        self.reg_head   = nn.Linear(d, 1)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        cls_token = torch.full((x.size(0), 1), self.cls_id, device=x.device)
        x = torch.cat([cls_token, x], dim=1)
        h = self.pos_enc(self.embedding(x))
        h = self.encoder(h)
        cls_out = h[:, 0, :]
        return self.cls_head(cls_out), self.reg_head(cls_out).squeeze(-1)

ESM-2 is the heavyweight, Meta AI’s 8-million-parameter protein language model, pre-trained on hundreds of millions of protein sequences, is loaded from HuggingFace and fine-tuned with only the last two transformer layers and the task heads as trainable parameters. The rest of the model stays frozen:

class ESM2Model(SequenceModel):
    def __init__(self, cfg: DictConfig) -> None:
        super().__init__()
        self.esm, self.alphabet = esm.pretrained.esm2_t6_8M_UR50D()
        # Freeze all layers except last two
        for name, param in self.esm.named_parameters():
            param.requires_grad = any(
                f"layers.{i}" in name for i in [4, 5]
            )
        hidden = self.esm.embed_dim
        self.cls_head = nn.Sequential(nn.Linear(hidden, 128), nn.ReLU(),
                                       nn.Linear(128, 3))
        self.reg_head = nn.Sequential(nn.Linear(hidden, 64),  nn.ReLU(),
                                       nn.Linear(64, 1))

Fine-tuning ESM-2 is strictly GPU-only — running it on CPU is technically possible but impractical for any meaningful dataset size. This constraint is enforced cleanly: the ESM-2 option is simply not available in the smoke sweep configuration.

Multi-task training with W&B logging

The training loop optimizes both tasks simultaneously with a weighted sum of losses:

total_loss = cls_loss_weight x CrossEntropy +reg_loss_weight x MSE

Default weights are cls=1.0 and reg=0.1, since cross-entropy and MSE live on very different scales and this ratio keeps the gradient signal balanced. Every epoch logs a rich set of metrics to Weights & Biases:

wandb.log({
    "train/loss": train_loss,
    "train/cls_loss": train_cls_loss,
    "train/reg_loss": train_reg_loss,
    "train/cls_acc": train_cls_acc,
    "train/hydro_r2": train_r2,
    "val/loss": val_loss,
    "val/cls_loss": val_cls_loss,
    "val/reg_loss": val_reg_loss,
    "val/cls_acc": val_cls_acc,
    "val/hydro_r2": val_r2,
    "lr": scheduler.get_last_lr()[0],
    "grad_norm": grad_norm,
    "epoch": epoch,
})

The learning rate schedule is configurable between CosineAnnealingLR and ReduceLROnPlateau, with early stopping on val/loss at configurable patience. Gradient norm logging is particularly useful for diagnosing instability in the Transformer, where exploding gradients can appear silently without proper monitoring.

Train vs. sweep: two distinct workflows

One of the more important design distinctions in this project is the clean separation between train and sweep. They look similar — both invoke the model, both log to W&B — but they serve fundamentally different purposes, use different entry points, and produce different kinds of output.

**make train-smoke / make train-full** execute a single, deterministic training run. Hyperparameters are read directly from configs/default.yaml, W&B is disabled for the smoke variant (to keep local iteration fast and uncluttered), and when the run finishes, the best checkpoint is written to checkpoints/best_model.pt. This is the workflow you use when you already know your configuration — for final model training, for debugging a specific architecture, or for validating a code change in isolation. Each run is fully reproducible given the same config and seed.

make train-smoke   # single run, W&B off, smoke data, <3 min on CPU
make train-full    # single run, W&B on, full OAS data, GPU recommended

**make sweep-cpu / make sweep-gpu** launch a W&B sweep agent — a fundamentally different control flow. Instead of running train.py directly, they invoke sweep.py, which registers a sweep with the W&B backend, then loops over multiple runs. In each run, W&B samples a hyperparameter configuration from the search space defined in the sweep YAML, injects it into the training loop, logs all metrics, and passes the results back to the sweep controller to inform the next sample. W&B is always enabled here: the sweep infrastructure requires it.

make sweep-cpu   # 5-run random search, smoke data, CPU, lstm + transformer only
make sweep-gpu   # 50-run Bayesian search, full OAS data, GPU, all 3 architectures

The two sweep configurations differ in more than just scale. sweep_cpu.yaml uses a random search strategy: each run draws hyperparameters independently, which is fast and unbiased but does not learn from previous results. This is appropriate for the smoke context where each run costs only seconds and the goal is simply to verify that the sweep machinery works:

# configs/sweep_cpu.yaml
method: random
metric:
  name: val/cls_acc
  goal: maximize
parameters:
  model_type:
    values: [lstm, transformer]   # ESM-2 excluded on CPU
  learning_rate:
    distribution: log_uniform_values
    min: 1e-4
    max: 1e-2
  dropout:
    values: [0.1, 0.2, 0.3]
  num_layers:
    values: [1, 2]
  embedding_dim:
    values: [32, 64]
run_count: 5

sweep_gpu.yaml switches to Bayesian optimisation. The W&B sweep controller fits a Gaussian process surrogate model over the observed (hyperparameter → metric) pairs from completed runs, and uses it to propose configurations in regions of parameter space that are most likely to improve on the current best. This is meaningfully more sample-efficient than random search when each trial is expensive — which it is at GPU scale with 500K sequences, a deep Transformer, and up to 50 training epochs per run:

# configs/sweep_gpu.yaml
method: bayes
metric:
  name: val/cls_acc
  goal: maximize
early_terminate:
  type: hyperband
  min_iter: 5
parameters:
  model_type:
    values: [lstm, transformer, esm2]   # ESM-2 included on GPU
  learning_rate:
    distribution: log_uniform_values
    min: 5e-5
    max: 5e-3
  batch_size:
    values: [128, 256, 512]
  dropout:
    distribution: uniform
    min: 0.05
    max: 0.4
  num_layers:
    values: [2, 3, 4, 6]
  d_model:
    values: [128, 256, 512]
  cls_loss_weight:
    values: [0.5, 1.0, 2.0]
run_count: 50

The GPU sweep also enables Hyperband early termination: runs that are underperforming after a minimum number of epochs are stopped automatically, freeing the GPU for more promising configurations. This is a practical necessity at 50 runs — without early termination, the sweep would be dominated by time spent on clearly bad configurations.

The practical consequence of this separation is that the two workflows serve entirely different phases of the project. Early in development, train-smoke gives fast, clean feedback on code changes. Once the architecture is stable, sweep-cpu confirms the sweep machinery works. On Kaggle, sweep-gpu does the real work of finding the optimal configuration across the full search space. And when a winner emerges from the sweep, a final train-full run with those hyperparameters committed to configs/default.yaml produces the production checkpoint. Each step is reproducible and independently verifiable — which is exactly what production-grade MLOps requires.

Running on Kaggle with GPU acceleration

Local development and testing happen on CPU with the smoke dataset. Production training — full OAS data, deep hyperparameter sweeps, ESM-2 fine-tuning — runs on Kaggle Notebooks with free GPU access. The workflow is deliberately simple and fully reproducible.

Open a Kaggle Notebook with GPU enabled (T4 x2 or P100), enable internet access, and add WANDB_API_KEY and HF_API_KEY to Kaggle Secrets. Then bootstrap the environment:

# Step 1: Clone the repo and install dependencies
!git clone https://github.com/Dima806/antibody_sequence_modeling_example.git
%cd antibody_sequence_modeling_example
!pip install -e ".[esm]" --quiet
!pip install -r requirements_kaggle.txt --quiet

With the environment ready, run a smoke test first to verify that everything wires up correctly before committing to a long GPU run:

# Step 2: Smoke run — validates the pipeline end-to-end on CPU in <3 minutes
!cd antibody_sequence_modeling_example && make data-smoke && make sweep-cpu

This generates the ≈2K synthetic dataset, validates it against the schema, and launches a 5-run random W&B sweep over BiLSTM and Transformer architectures on smoke data. If this completes cleanly, you have confidence that the full run won’t fail on a configuration error three hours in.

Then trigger the real training:

# Step 3: Full GPU run — downloads OAS data and launches 50-run Bayesian sweep
!cd antibody_sequence_modeling_example && make data-full && make sweep-gpu

make data-full downloads the pre-processed OAS snapshot from Zenodo (no authentication required), validates it, and writes it to data/full/sequences_full.csv. make sweep-gpu then launches the Bayesian W&B sweep over 50 runs, searching across model type, learning rate, batch size, dropout, number of layers, and model dimensionality simultaneously. The Bayesian optimization strategy (as opposed to random search on CPU) is justified here: with 50 runs on a GPU, the cost of each trial is high enough that it is worth using the information from previous runs to guide the search.

Experiment tracking and results sharing via W&B

Every training run — smoke or full, CPU or GPU — logs automatically to the W&B project. This is where the real value of integrating W&B from the start becomes visible.

The W&B sweep dashboard shows all 50 GPU runs side by side, with parallel coordinates plots revealing which hyperparameter combinations drove the best validation accuracy and R2. Run comparison tables make it easy to identify whether, for example, the Transformer consistently outperforms the BiLSTM at a given parameter budget, or whether the learning rate is the dominant factor at small batch sizes.

Sharing results with collaborators is a single step: add them to the W&B team, and they have full access to all run histories, metric curves, model checkpoints, and sweep analyses — without needing to clone the repo, install dependencies, or re-run anything. In a context where computational teams and scientists need to communicate about model confidence and prediction quality, this kind of frictionless sharing is important.

The Makefile as operational interface

One detail worth highlighting explicitly is the role of the Makefile as the project's single operational interface. Every meaningful action — from setup to testing to training to sweeping — is a make target:

setup:          ## Install uv, package (editable), pre-commit hooks
data-smoke:     ## Generate ~2K synthetic CDR-H3 sequences + validate
data-full:      ## Download OAS dataset from Zenodo + validate
train-smoke:    ## Smoke training run, W&B disabled
train-full:     ## Full training run on OAS data
sweep-cpu:      ## Random W&B sweep, 5 runs, CPU, smoke data
sweep-gpu:      ## Bayesian W&B sweep, 50 runs, GPU, full data
test:           ## pytest + coverage
lint:           ## pre-commit (ruff lint + format, YAML/TOML, hooks)

This means that documentation about how to run the project never drifts out of sync with the actual commands, because the documentation is the commands. It also means that Kaggle Notebook cells, CI pipelines, and local development all speak the same language — there are no environment-specific scripts to maintain separately.

Final words and pipeline limitations

The immediate output is a model that can score a library of CDR-H3 sequences for length class and hydrophobicity. Candidates with high predicted hydrophobicity (a proxy for aggregation risk) or anomalous loop lengths can be filtered automatically, focusing experimental effort on the developable fraction of the library.

But it is worth being clear about what this project is and is not. It is not a deployment-ready inference service. There is no REST API, no containerised model server, no automated retraining trigger. The training pipeline requires a human to open a Kaggle notebook and run cells. For an individual researcher or a small team working in a resource-constrained environment, that is often entirely fine — the bottleneck is usually experiment iteration, not deployment automation. But if you are building toward a system where a pipeline automatically calls a model endpoint, the pieces that are missing here are well-defined and addable: wrap the best checkpoint in FastAPI, containerise with Docker, and deploy to any cloud. The model and the training code are ready for that.

Having said that, the MLOps patterns here — dual-environment training with an explicit manual bridge, schema-validated data pipelines, Makefile-driven workflows, W&B sweep integration, uv-managed dependencies — apply to any ML project that needs to move cleanly from exploratory development to reproducible GPU training, in an environment where cloud compute budgets are constrained.

Drop your questions in the comments below 😊