The Bear's Toes https://bear-toes.pages.dev/ quarto-1.9.38 Sun, 13 Apr 2025 23:00:00 GMT Distributed Training with DistributedDataParallel https://bear-toes.pages.dev/posts/torch-distributed-explained/

# What is DistributedDataParallel (DDP)?

DistributedDataParallel is a way to parallelize training across multiple GPUs or nodes. It is an extension of DataParallel that provides more flexibility and scalability. DataParallel (DP) is an older approach to data parallelism. DP is trivially simple (with just one extra line of code) but it is less performant. DDP improves upon the architecture in a few ways:

DataParallel DistributedDataParallel
Simpler to use More involved changes to use
More overhead; model is replicated and destroyed at each forward pass Model is replicated only once at the start
Only supports single-node parallelism Supports single-node and multi-node parallelism
Slower; uses multithreading on a single process and runs into Global Interpreter Lock (GIL) contention Faster (no GIL contention) because it uses multiprocessing

# Multi-GPU Training with DDP

DDP uses multiprocessing to copy the model to each GPU (rank). This allows the model (and code) to only be copied to each process once at the start of the script. Multiprocessing pickles Python objects to serialize across processes. This means all objects must be pickleable.

ImportantAll objects sent to each process must be pickleable

What can be pickled and unpickled?

Some objects that can’t be pickled:

  • Generators
  • Database connections
  • Sockets
  • File descriptors
  • Lambdas

The basic outline of DDP training is:

  1. Setup the communications by setting the host and port
  2. Spawn a training process for each GPU with torch.multiprocessing.spawn
  3. Initialize the process group using init_process_group:
    • GPU - "nccl"
    • CPU - "gloo"
  4. Wrap the model with DistributedDataParallel
  5. Create a DistributedSampler and DataLoader for the dataset
  6. Train the model and update sampler with the epoch
  7. Destroy the process group using destroy_process_group
Code 
from torch.utils.data.distributed import DistributedSampler
from torch.distributed import init_process_group, destroy_process_group

def main(
    rank: int, # rank is the GPU number
    world_size: int, # world_size is the number of processes, typically set to the number of GPUs
    train_path: str,
    random_state: int,
    lr: float,
    epochs: int,
    num_workers: int,
    batch_size: int,
) -> None:
    # ... Other setup

    # initialize the process group for distributed training
    init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
        rank=rank,
        world_size=world_size,
    )

    # we need to divide the workers and batch across the different processes used in distributed training
    num_workers_per_proc = num_workers // world_size # avoids CPU contentionn
    batch_size_per_proc = batch_size // world_size   # avoids OOM

    # DistributedSampler ensures that training data is chunked across GPUs without overlapping samples
    train_sampler = DistributedSampler(train_dataset)
    val_sampler = DistributedSampler(
        val_dataset,
        shuffle=False,  # don't shuffle the validation dataset
        drop_last=True, # DistributedSampler will append additional samples to fill an incomplete batch.  We don't want that for the validation dataset.
    )

    train_dataloader = DataLoader(
        train_dataset,
        shuffle=False,  # don't shuffle if using DistributedSampler as that's done within the sampler
        sampler=train_sampler,
        num_workers=num_workers_per_proc,
        batch_size=batch_size_per_proc,
        pin_memory=True,
        collate_fn=collate_rowgroups,
    )
    val_dataloader = DataLoader(
        val_dataset,
        shuffle=False,
        sampler=val_sampler,
        num_workers=num_workers_per_proc,
        batch_size=batch_size_per_proc,
        pin_memory=True,
        collate_fn=collate_rowgroups,
    )

    # set up the NN model as normal and then wrap with DDP
    model.to(rank)
    model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    for epch in range(0, epochs):
        # need to call `set_epoch()` at the beginning of each epoch before creating the
        # `DataLoader` iterator to make shuffling work properly across multiple epochs
        train_sampler.set_epoch(epch)

        # ... training loop

    # ... diagnostics

    # cleanly shutdown distributed processes
    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    import os
    import torch

    os.environ["MASTER_ADDR"] = "localhost" # for single node on local compute
    os.environ["MASTER_PORT"] = "12345" # any free port

    world_size = torch.cuda.device_count()  # number of GPUs

    # ... CLI args parsing

    # spawn multiple processes equal to world_size first argument passed in will be the rank
    torch.multiprocessing.spawn(
        main,
        args=(
            world_size,
            args.train_path,
            args.random_state,
            args.lr,
            args.epochs,
            args.num_workers,
            args.batch_size,
        ),
        nprocs=world_size,  # this is used to set the `rank` parameter.  It is passed as the first argument
    )

1. Communication: host and port

Setting up distributed training on a single (local) node is as simple as setting the host and port as below. To setup multi-node see torchrun.

Code 
import os
os.environ["MASTER_ADDR"] = "localhost" # for single node on local compute
os.environ["MASTER_PORT"] = "12345" # any free port

2. Spawn a process on each rank (GPU)

Since torch.multiprocessing follows the same API as multiprocessing. To spawn a new process we pass the function to run, the arguments as a tuple, and specify the number of processes (usually the number of GPUs).

Code 
import torch

# spawn multiple processes equal to world_size first argument passed in will be the rank
torch.multiprocessing.spawn(
    main,
    args=(
        world_size,
        args.train_path,
        args.random_state,
        args.lr,
        args.epochs,
        args.num_workers,
        args.batch_size,
    ),
    nprocs=world_size,  # this is used to set the `rank` parameter.  It is passed as the first argument
)

3. Constructing the process group

  • First, before initializing the group process, call set_device, which sets the default GPU for each process. This is important to prevent hangs or excessive memory utilization on GPU:0
  • The process group can be initialized by TCP (default) or from a shared file-system. Read more on process group initialization.
  • init_process_group initializes the distributed process group.
  • Read more about choosing a DDP backend.
Code 
from torch.distributed import init_process_group

# initialize the process group for distributed training
init_process_group(
    backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
    rank=rank,
    world_size=world_size,
)

5. Constructing the DDP model

  • device_ids - 1) For single-device modules, device_ids can contain exactly one device id, which represents the only CUDA device where the input module corresponding to this process resides. Alternatively, device_ids can also be None. 2) For multi-device modules and CPU modules, device_ids must be None. (From the DDP docs)
model.to(rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

4. Distributing the data with DistributedSampler

Dividing the workload

  • DistributedSampler chunks the input data across all distributed processes, without overlap. If we have 4 GPUs then each process will only load 1/4 of the training dataset.
  • The batch_size needs to be divided among the processes (GPUs). Each process will receive an input batch of batch_size_per_proc; the effective batch size is batch_size_per_proc * world_size, if the batch_size is 64 and world_size is 4 GPUs, then the effective batch size is still 64 in total.
  • The num_workers also needs to be divided among the processes (GPUs). Each proces will receive num_workers_per_proc.
Cautionbatch_size and OOM

If the batch_size isn’t divided among the processes then then each process gets a full batch and the effective batch size is now xworld_size larger and we are likely to run out of CPU or GPU memory if not careful.

Code 
# we need to divide the workers and batch across the different processes used in distributed training
num_workers_per_proc = num_workers // world_size # avoids CPU contention
batch_size_per_proc = batch_size // world_size   # avoids OOM

Setting Up the DistributedSampler

  • shuffle - by default, the DistributedSampler will shuffle the dataset. We don’t want to shuffle the validation dataset.
  • drop_last - by default, the DistributedSampler will append additional samples to fill an incomplete batch (e.g. there’s 100 training samples with batch_size=64 there would be one batch of 36 samples). We don’t want to repeat samples for the validation dataset as that would change the metrics.
  • pin_memory - For large datasets that are loaded into the CPU in the Dataset, pinning the memory can speed up the host to device transfer (see this NVIDIA blog for more details).
Code 
# DistributedSampler ensures that training data is chunked across GPUs without overlapping samples
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(
    val_dataset,
    shuffle=False,  # don't shuffle the validation dataset
    drop_last=True, # DistributedSampler will append additional samples to fill an incomplete batch.  We don't want that for the validation dataset.
)

train_dataloader = DataLoader(
    train_dataset,
    shuffle=False,  # don't shuffle if using DistributedSampler as that's done within the sampler
    sampler=train_sampler,
    num_workers=num_workers_per_proc,
    batch_size=batch_size_per_proc,
    pin_memory=True,
    collate_fn=collate_rowgroups,
)
val_dataloader = DataLoader(
    val_dataset,
    shuffle=False,
    sampler=val_sampler,
    num_workers=num_workers_per_proc,
    batch_size=batch_size_per_proc,
    pin_memory=True,
    collate_fn=collate_rowgroups,
)

7. Shuffling across epochs

  • Calling the set_epoch() method on the DistributedSampler at the beginning of each epoch is necessary to make shuffling work properly across multiple epochs. Otherwise, the same ordering will be used in each epoch.
Code 
# ... Neural Network setup

for epch in range(0, epochs):
    # need to call `set_epoch()` at the beginning of each epoch before creating the
    # `DataLoader` iterator to make shuffling work properly across multiple epochs
    train_sampler.set_epoch(epch)

    # ... training loop

6. Running the distributed training job

  • rank is auto-allocated by DDP when calling torch.multiprocessing.spawn.
  • world_size is the number of processes across the training job. For GPU training, this corresponds to the number of GPUs in use, and each process works on a dedicated GPU.
  • Both rank and world_size are new parameters to main(). Because of how spawning processes works, rank needs to be the first parameter to the calling function, main(rank, ...).
TipPyTorch Multiprocessing

PyTorch’s torch.multiprocessing package is a wrapper around the native multiprocessing module and the API is 100% compatible.

Code 
import torch

# ... CLI args parsing

world_size = torch.cuda.device_count()  # number of GPUs

torch.multiprocessing.spawn(
    main,
    args=(
        world_size,
        args.train_path,
        args.random_state,
        args.lr,
        args.epochs,
        args.num_workers,
        args.batch_size,
    ),
    nprocs=world_size,  # this is used to set the `rank` parameter.  It is passed as the first argument
)

MLFlow logging

Since there are now multiple processes runnning the same code, the same logging will happen on each process. MLFlow doesn’t know how to distinguish that there are different processes logging the same metric. We can guard against this by only logging on the main process (GPU0):

Code 
import mlflow

if rank == 0:
    mlflow.log_metrics(
        {
            "train loss": train_loss,
            "train accuracy": train_accuracy,
            "val loss": val_loss,
            "val accuracy": val_accuracy,
        },
        step=epoch,
    )

# Gradients, Losses, and Metrics

Under the hood, DDP synchronizes and gathers the gradients across all processes. However, any other ad-hoc value calculated in your code is not; e.g. losses and metrics.

NoteGradients are synchronized

Model gradients are synchronized across processes during the backward pass. This means that the model in each process is the same! See DDP: Internal Design.

CautionLosses are not synchronized

Even though the model is the same in each process, the loss is calculated on only the portion of the batch that each process sees. The losses don’t need to be synchronized for training but we may want to synchronize the losses for logging or definitelty when calculating metrics on the hold-out (validation) dataset.

We could avoid this by not using a DistributedSampler for the validation set, but then only 1 process would be used to calculate the loss for the whole validation set each epoch, which will be slow.

So if each process is calculating and accumulating losses and metrics separately, how do we log those and report as if there were a single process? Well, those values will need to be gathered and then accumulated. Say we have 4 processes, one for each GPU, and each is processing 1/4 of the training dataset. We want to report the loss for each epoch. If we log the loss in each process, we will have 4 different losses. We can gather and combine them in a few ways. Since the loss is just a number value we can use torch.distributed.reduce or torch.distributed.all_reduce:

Example: torch.distributed.reduce

In this example, we gather and combine using summation with the dist.ReduceOp.SUM, all the loss_tensors into the process 0 tensor (dst=0). Each tensor in each process must be the same shape. Since we are assigning the values in each processes’s loss_tensor to it’s rank, we expect the final gathered values to be 0 + 1 + 2 + 3 = 6 in the main process loss_tensor.

Code 
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.distributed import init_process_group

def reduce_tensor(rank: int, world_size: int) -> None:
    init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
        rank=rank,
        world_size=world_size,
    )
    torch.cuda.set_device(rank) # tell each device (GPU) which one it is.

    loss_tensor = torch.tensor([rank, rank]).cuda()
    print(loss_tensor)

    dist.reduce(loss_tensor, op=dist.ReduceOp.SUM, dst=0, async_op=True)
    print(loss_tensor)

    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"

    num_gpu = torch.cuda.device_count()
    mp.spawn(reduce_tensor, nprocs=num_gpu, args=(num_gpu,))
tensor([1, 1], device='cuda:1')
tensor([3, 3], device='cuda:3')
tensor([2, 2], device='cuda:2')
tensor([0, 0], device='cuda:0')

tensor([6, 6], device='cuda:0')
tensor([1, 1], device='cuda:1')
tensor([2, 2], device='cuda:2')
tensor([3, 3], device='cuda:3')

Example: torch.distributed.all_reduce

In this example, we gather and combine using summation with the dist.ReduceOp.SUM, all the loss_tensors into all the processes. Each tensor in each process must be the same shape. Since we are assigning the values in each processes’s loss_tensor to it’s rank, we expect the final gathered values to be 0 + 1 + 2 + 3 = 6 in the all the processes’s loss_tensor.

Code 
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.distributed import init_process_group

def all_reduce_tensor(rank: int, world_size: int) -> None:
    init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
        rank=rank,
        world_size=world_size,
    )
    torch.cuda.set_device(rank) # tell each device (GPU) which one it is.

    loss_tensor = torch.tensor([rank, rank]).cuda()
    print(loss_tensor)

    dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM, dst=0, async_op=True)
    print(loss_tensor)

    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"

    num_gpu = torch.cuda.device_count()
    mp.spawn(all_reduce_tensor, nprocs=num_gpu, args=(num_gpu,))
tensor([2, 2], device='cuda:2')
tensor([3, 3], device='cuda:3')
tensor([1, 1], device='cuda:1')
tensor([0, 0], device='cuda:0')

tensor([6, 6], device='cuda:2')
tensor([6, 6], device='cuda:3')
tensor([6, 6], device='cuda:1')
tensor([6, 6], device='cuda:0')

Example: torch.distributed.gather_object

Syncing across processing is simple enough for tensors, but if we have a number of values to gather (say a bunch of metrics for example) it would be easier to only need to gather once and store those values in an appropriate data structure. Most of the distributed gathering function only work on tensors, but we can use torch.distributed.gather_object and / or torch.distributed.all_gather_object to pass pickleable Python objects between ranks.

In this example, we want to track the losses and number of samples in a Counter so that we can combine and calculate the mean loss after gathering. We gather each loss counter to rank 0. Each counter is placed into the gather_list, which must have all elements set to None initially. When calling, dist.gather_object, the gather_list must only exist in the rank being gathered to (dst=0 or rank 0 in this case). Then we use functools.reduce to sum all the Counters gathered in the gather_list.

Code 
import os
import operator
from collections import Counter
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.distributed import init_process_group

def gather_object(rank: int, world_size: int) -> None:
    init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
        rank=rank,
        world_size=world_size,
    )
    torch.cuda.set_device(rank) # tell each device (GPU) which one it is.

    losses = Counter(
        {"loss": 0.01, "num_samples": rank}
    )
    print(losses)

    gather_list = [None for _ in range(world_size)]

    dist.gather_object(losses, gather_list if rank == 0 else None, dst=0)
    losses = functools.reduce(operator.add, gather_list)

    print(losses)
    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"

    num_gpu = torch.cuda.device_count()
    mp.spawn(gather_object, nprocs=num_gpu, args=(num_gpu,))
Counter({'num_samples': 1, 'loss': 0.01})
Counter({'num_samples': 2, 'loss': 0.01})
Counter({'num_samples': 4, 'loss': 0.01})
Counter({'num_samples': 3, 'loss': 0.01})

Counter({'num_samples': 10, 'loss': 0.04})
Counter({'num_samples': 4, 'loss': 0.01})
Counter({'num_samples': 3, 'loss': 0.01})
Counter({'num_samples': 2, 'loss': 0.01})

Example: torch.distributed.all_gather_object

In this example, we want to track the losses and number of samples in a Counter so that we can combine and calculate the mean loss after gathering. We gather each loss counter to each rank. Each counter is placed into the gather_list, which must have all elements set to None initially. Then we use functools.reduce to sum all the Counters gathered in the gather_list.

Code 
import os
import operator
from collections import Counter
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from torch.distributed import init_process_group

def all_gather_object(rank: int, world_size: int) -> None:
    init_process_group(
        backend="nccl" if torch.cuda.is_available() else "gloo",  # CPU only works on gloo backend
        rank=rank,
        world_size=world_size,
    )
    torch.cuda.set_device(rank) # tell each device (GPU) which one it is.

    losses = Counter(
        {"loss": 0.01, "num_samples": rank}
    )
    print(losses)

    gather_list = [None for _ in range(world_size)]

    dist.all_gather_object(gather_list, losses)
    losses = functools.reduce(operator.add, gather_list)

    print(losses)
    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"

    num_gpu = torch.cuda.device_count()
    mp.spawn(all_gather_object, nprocs=num_gpu, args=(num_gpu,))
Counter({'num_samples': 1, 'loss': 0.01})
Counter({'num_samples': 4, 'loss': 0.01})
Counter({'num_samples': 2, 'loss': 0.01})
Counter({'num_samples': 3, 'loss': 0.01})

Counter({'num_samples': 10, 'loss': 0.04})
Counter({'num_samples': 10, 'loss': 0.04})
Counter({'num_samples': 10, 'loss': 0.04})
Counter({'num_samples': 10, 'loss': 0.04})

# Multi-Node Training

Now that we’ve setup DDP, we have the option of running on multiple nodes. There are a few extra bits that need to be changed before we run multinode training:

  • change a few lines to enable torchrun
  • install the project as a python package for torchrun to work properly
  • add fault tolerance

# Resources

  1. DistributedDataParallel - Documentation for DDP.
  2. Getting Started with Distributed Data Parallel - Good starting point to understand DDP and writing a training script using DDP for single-node and multi-node.
  3. Writing Distributed Applications with Pytorch - In-depth article about writing distributed applications in PyTorch and how communication works under the hood.
  4. DistributedSampler - Used in conjunction with a DataLoader, the DistributedSampler enables each process to only load the data it processes, rather than all the data to be processed.
  5. DistributedDataParallel training in PyTorch - Explaination of how DDP works and how to use it.
  6. (AML) Distributed GPU training guide (SDK v2)
  7. PyTorch DistributedDataParallel Example In Azure ML - Multi-Node Multi-GPU Distributed Training - Example of DDP for AML.
  8. PyTorch Distributed: Experiences on Accelerating Data Parallel Training - Paper on how Facebook designed DDP to be faster for distributed training.
]]> deep learning pytorch distributed systems https://bear-toes.pages.dev/posts/torch-distributed-explained/ Sun, 13 Apr 2025 23:00:00 GMT Statistics on Stream Data https://bear-toes.pages.dev/posts/stream-statistics/

# Statistics on Stream Data

Simple statistics are common in a number of software and data domains because they are simple to calculate and simple to understand. Most common are the arithmetic mean and variance. For a given set of values, , the mean, , is a measure of central tendency and locates a central value of the distribution. The variance, , provides a measure of the dispersion of the observed values of about the central value (Schott 2016, 20–25).

Suppose we have a set of values, ; then:

Code 
import torch

X = torch.tensor([1.0, 2.0, 1.0, 2.0, 4.0, 5.0])

print(f"mean = {X.mean()}")
print(f"var  = {X.var(correction=0)}")
print(f"sample_var = {X.var()}")
mean = 2.5
var  = 2.25
sample_var = 2.700000047683716

This should be no surprise if you’ve worked with data before and in theory, the mean and variance are straightforward to calculate. But in practice, datasets nowadays have become very large in the age of “Big Data”. Not only are datasets now too big to fit into memory, datasets are ever growing and constantly being updated. How do we calculate these measures if the data is too large to process in memory? How can we update these measures when new data is added to the dataset? Do we have to reprocess every value? Storing large datasets can become costly but reprocessing large datasets just to recalculate a simple statistic can be prohibitive. Many datasets have a temporal component and simple statistics can be calculated on a moving or sliding window because the most important data is the most recent data. Many other datasets can’t just calculate over a lsiding window. For example, deep learning models are now being trained on very large datasets, i.e. the whole of the internet.Since the entire internet can’t fit into memory of even the largest compute instance, these models are trained in batches. Typically, the mean is computed over a batch (BatchNorm) but when evaluating the model on a hold-out dataset, we want to compute statistics over the whole of the hold-out dataset.

We could just keep a running sum of the numerator and running count of the denominator but there are some drawbacks. Calculating the numerator of the mean could result in an arithmetic overflow when dealing with large numbers. Calculating the numerator of the variance involves recalculating the mean for each new batch of data and involves calculating the sums of squares, which can lead to numerical instability.

Suppose our dataset arrives in batches of 2: . We can update the mean and sample variance as:

TipUse the sample variance

Since most datasets don’t truely sample the whole population and because each batch is effectively a sample, we’ll want to compute the sample variance, , rather than the population variance.

where:

  • is the running count of all values seen so far
  • is the count of the values in the new batch of data
  • is the mean of all values seen so far
  • is the mean of the new batch of data
  • is the sample variance of all the values seen so far
  • is the sample variance of the new batch of data

The updated mean is a linear combination of the mean over the seen data and the mean over new data. The updated variance is a linear combination of the seen data variance and new data variance plus a correction by the means.

Code 
import torch

def update_mean(a: int, b: int, mean_a: float, mean_b: float) -> float:
    return (a * mean_a + b * mean_b) / (a + b)

def update_variance(a: int, b: int, mean_a: float, mean_b: float, var_a: float, var_b: float) -> float:
    
    within_sample_var = (a - 1) * var_a + (b - 1) * var_b
    between_sample_var = (a*b) / (a+b) * (mean_a - mean_b)**2
    
    return (within_sample_var + between_sample_var) / (a+b-1)

X = torch.tensor([[1.0, 2.0], [1.0, 2.0], [4.0, 5.0]])

# while not strictly correct, we can initialize 
# the mean and variance to 0 rather than NaN
count = 0
mean = 0.0
var = 0.0
for x in X:
    mean_updated = update_mean(count, x.numel(), mean, x.mean())
    var_updated = update_variance(count, x.numel(), mean, x.mean(), var, x.var())
    
    count += x.numel()
    mean = mean_updated
    var = var_updated

assert mean == X.mean()
assert var == X.var()

print(f"mean = {mean}")
print(f"var  = {var}")
mean = 2.5
var  = 2.700000047683716

Derivation

Batch Update Mean

We can think of the two batches we are combining as two ranges that are being combined. So if batch A is between and batch B is between and we are starting with batch A and adding batch B, the combined range is then , where the total number of samples observed is .

Knowing this we can express the combined mean as:

Batch Update Variance

For the variance, we’ll first derive the population variance and then work out the sample variance.

Let’s start by simplifying the variance equation.

Now, we can follow a similar process as we did when working out the combined mean, starting from equation 5.

Now, to find the sample variance, , we can use Bessel’s correction, :

then we have:

we can then substitute these in to find the sample variance,

Excersises Left Up to the Reader

Now that you know how to derive batch updates for the mean and variance, you should be well-equiped to derive the covariance, , and pearson correlation, :

Further Reading

References

Schott, James R. 2016. Matrix Analysis for Statistics. 3rd ed. Wiley.
]]> probability machine learning streaming https://bear-toes.pages.dev/posts/stream-statistics/ Sat, 12 Apr 2025 23:00:00 GMT Creating Memory Leaks with Parquet Metadata https://bear-toes.pages.dev/posts/parquet_metadata/

# The Setup

In my day job, I develop computer vision models to detect methane emissions in satellite images. We have a training dataset of a few million chips (a subset of the full image from the satellite) that is about 2.5TB on disk stored in a couple thousand parquet files. Now, there’s not a lot of real labelled data for us to use, so a big part of my work is how to generate synthetic methane data.

Recently, our training jobs started running out of CPU memory. Each training epoch takes about 30 minutes, so we try to saturate the GPUs with large batches to train faster. The main bottleneck is data loading and keeping the GPUs saturated with data. To saturate the GPUs, we have multiple workers prefetching batches so the GPUs have high utilization. This means that the CPU needs to hold many extra batches in memory ready and waiting. Suffice to say, our training jobs run with fairly high memory usage; a small change to our training script can OOM (Out of Memory) the whole job.

# The Situation

Recently, our training jobs started OOMing in the first epoch, running out of CPU memory. We tried adjusting the obvious levers: batch size, the number of workers, and prefetching but the jobs were still running out of memory.

O.K. so what changed? Our training data had changed in two ways: 1. We started tracking additional metadata about our generated training dataset, going from ~15 columns to ~150 columns. The additional columns were mostly float64 data type with a few string or list of string columns. 2. The row group size in the parquet files decreased from 10 to 1, meaning, each row group has a single row.

The new columns aren’t loaded during the training loop, so this couldn’t be the cause of the memory explosion.

We decreased the row group size to 1 so that random sampling of training examples could be truly random. For performance reasons, we were randomly sampling row groups because Parquet does not support partial row group reads. It’s much faster to use the whole row group as a “sample” rather than an individual row because it avoids additional IO from loading the same row group multiple times. With row groups of 10 rows, we used all 10 rows, but then the model always saw those 10 rows together. This is actually even worse because due to how the data is generated, all the samples in a file are from the same MGRS tile. So the rows in a row group are correlated. With row groups of 1, we can still sample by row groups but now batches are composed of truly random samples.

We had benchmarked data loading with row groups of 10 and row groups of 1 and found them comparable. So why not switch and have better randomization?

# Finding the Source

# Experimenting

Let’s inspect the metadata for a file of 200 rows, 15 columns, and row group size of 10.

Code 
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

num_cols = 15
num_rows = 1_000
data = {f"{x}": np.random.rand(num_rows) for x in range(num_cols)}
df = pd.DataFrame(data)
df.head()

file_path = "data.parquet"
df.to_parquet(
    file_path,
    compression="zstd",
    compression_level=9,
    row_group_size=10,
)

file = pq.ParquetFile(file_path)
print(file.metadata)
<pyarrow._parquet.FileMetaData object at 0x103c73ec0>
  created_by: parquet-cpp-arrow version 23.0.1
  num_columns: 15
  num_rows: 1000
  num_row_groups: 100
  format_version: 2.6
  serialized_size: 163754


The metadata for this parquet file is 14,895 bytes or ~0.14 MB. Now let’s see a file of 150 columns with row group size of 1.

Code 
num_cols = 150
num_rows = 1_000
data = {f"{x}": np.random.rand(num_rows) for x in range(num_cols)}
df = pd.DataFrame(data)
df.head()

file_path = "data.parquet"
df.to_parquet(
    file_path,
    compression="zstd",
    compression_level=9,
    row_group_size=1,
)

file = pq.ParquetFile(file_path)
print(file.metadata)
<pyarrow._parquet.FileMetaData object at 0x1088eeb10>
  created_by: parquet-cpp-arrow version 23.0.1
  num_columns: 150
  num_rows: 1000
  num_row_groups: 1000
  format_version: 2.6
  serialized_size: 16139711


The metadata for this parquet file is 14,642,267 or 14MB, that’s massive for just metadata!

O.K. we’ve discovered that the metadata of our Parquet files had increased, but that should only be held in memory when reading from the file. To randomly sample from our dataset, we create an mapping of the index to the file and row group and then randomly sample the index. Each time a ParquetFile is initialized, it reads the file metadata into memory. This is expensive to do each time we read from the file, so we cache the object.

Code 
import functools
from dataclasses import dataclass

from torch.utils.data import Dataset

@dataclass
class Partition:
    """Store information about the partition for faster retrieval."""

    file: str
    row_group: int


@functools.cache
def create_parquet_file(file: str) -> pq.ParquetFile:
    """Caches the creation of the ParquetFile."""
    return pq.ParquetFile(file)


class ArrowDataset(Dataset):

    def __init__(self, dataset: ds.FileSystemDataset) -> None:
    self.dataset = dataset
        self.partitions: list[Partition] = self.create_partition_mapping(self.dataset)

    def create_partition_mapping(self, dataset: ds.Dataset) -> list[Partition]:
        """Create a mapping of partition number to the file and row_group."""
        partitions = []

        for file in dataset.files:
            parquet_file = create_parquet_file(file)

            for row_group in range(parquet_file.num_row_groups):
                partitions.append(Partition(file, row_group))

        return partitions


    def __len__(self) -> int:
        """Return the number of row_groups in the dataset.  This should correspond to RecordBatches."""
        return len(self.partitions)


Our training data is roughly 2,000 files with 1,000 rows per file. At 14MB per file, that’s 27GB of memory just for metadata! But it gets worse! Since we train with 4 GPUs and have 2 workers per GPU and each of those workers has its own process, that means we’re holding roughly 27GB * 4 * 2 = 216GB of metadata in memory! That’s atrocious!

As it turns out, there are three variables that contribute to most of the metadata size (the number of rows being constant): the row group size, the number of columns, and the column name length.

Code 
from pathlib import Path

import altair as alt

num_rows = 1_000
num_columns = []
num_rg = []
col_name_lengths = []
metadata_size = []

for num_cols in [1, 10, 50, 100, 150]:
    for name_length in [1, 10, 50, 100]:
        col_name = "x" *  name_length
        data = {f"{x}{col_name}": np.random.rand(num_rows) for x in range(num_cols)}
        df = pd.DataFrame(data)

        for row_group_size in [1_000, 100, 10, 1]:
            path = f"file_{num_cols}cols_{row_group_size}rgsize_{name_length}namelen.parquet"
            df.to_parquet(
                path,
                compression="zstd",
                compression_level=9,
                row_group_size=row_group_size,
            )
            file = pq.ParquetFile(path)

            num_columns.append(num_cols)
            num_rg.append(row_group_size)
            metadata_size.append(file.metadata.serialized_size)
            col_name_lengths.append(name_length)

df = pd.DataFrame({
    "num_cols": num_columns,
    "row_group_size": num_rg,
    "col_name_lengths": col_name_lengths,
    "metadata_size_b": metadata_size,
}).assign(
    metadata_size_mb=lambda df: df.metadata_size_b / 1024**2,
)


When adding the new columns to our dataset, we used long descriptive names which turns out to also be a contributing factor.

TODO: talk about chart

# Removing the Memory Landmines

Since statistics are calculated for each row group, increasing the number of row groups also increases the amount of statistics calculated and stored in the metadata. The statistics are used during predicate push down to check if the row group matches the predicate and if not, the row group can be skipped. Smaller row groups reduces the utility of the statistics, and storing statistics for a single row is just silly. So we can remove the statistics for our purposes.

Row groups are also encoded (actually the column chunks in a row group are encoded). Encoding small row groups doesn’t allow for great compression and for our use case, encoding the row groups increased the size of the metadata. So we can remove the encoding.

Parquet files created with Pandas or PyArrow, by default, include the Arrow schema as well. The Arrow schema includes additional information. to faithfully recreate the original Arrow data. We aren’t using complex types that would benefit from this, so we can remove the extra schema.

Shrinking the size of the metadata boils down to three lines when writing the parquet files:

Code 
num_cols = 150
num_rows = 1_000
data = {f"{x}": np.random.rand(num_rows) for x in range(num_cols)}
df = pd.DataFrame(data)
df.head()

file_path = "data.parquet"
df.to_parquet(
    file_path,
    compression="zstd",
    compression_level=9,
    row_group_size=1,
    write_statistics=False,
    store_schema=False,
    use_dictionary=False
)

file = pq.ParquetFile(file_path)
print(file.metadata)
<pyarrow._parquet.FileMetaData object at 0x108d7d170>
  created_by: parquet-cpp-arrow version 23.0.1
  num_columns: 150
  num_rows: 1000
  num_row_groups: 1000
  format_version: 2.6
  serialized_size: 6934573


The size of the metadata is now 5,737,509 bytes or 5.4MB, that’s 40% of the initial 14MB.

And we can rerun the experiment with the new parameters.


So we were able to reduce the metadata size by 60% just by changing some parameters when writing the file. From the chart, we can see that by reducing the average column name length from 50 characters to 20 characters, we can shave off another few MB per file. That’s simple enough.

We can shave another another few MB off by reducing the number of columns. That’s not so simple. We want to store that extra data with our training data, it’s useful for analysis and debugging. Parquet isn’t great for storing wide data (lots of columns) but 150 columns doesn’t seem that wide. We don’t want to separate the extra data into separate files because that requires tracking two datasets and joining them; it’s doable but it’s extra work. We could also reduce the extra columns by grouping them all into a single column of JSON or pyarrow structs. This works really well for reducing the size of the metadata but it necessitates knowing exactly which columns to group together and this is variable as our generated training data is updated regularly and has a few configurations. So that would also be a bit of a headache.

# You’re Holding it Wrong

Now, you might be thinking, “Parquet isn’t intended for these workloads, why not use a more appropriate file format?” And to that I say, yes it isn’t a typical workload but I challenge you to find a better file format that is efficient for random access patterns and still support projection pushdowns. Our use case should be fairly common for training machine learning models. There is already a Pull Request in the Arrow project to make Parquet metadata reading faster for random access patterns.

The solution we’re now using, Parquet files without metadata nor encoding isn’t all that different from the Arrow format (A.K.A. Feather A.K.A. Arrow I.P.C.). We might consider storing our training data in Arrow format as PyArrow can directly read the bytes on disk without any additional processing. In theory, this could be faster than Parquet for our workloads but our initial benchmarks indicated that storing in Arrow format was no faster than storing in Parquet format.

We also tried selecting rows directly with pa.dataset.take() as that API is a bit cleaner and seems like the preferred way to select rows from a dataset. However, this was roughly 30% slower than rolling our solution.

# Resources

]]> parquet https://bear-toes.pages.dev/posts/parquet_metadata/ Sat, 01 Mar 2025 00:00:00 GMT Career Tips Learned the Hard Way https://bear-toes.pages.dev/posts/career_tips/

# Introduction

  • Getting a Job
    • Resume Tips
    • Networking
    • Cover Letter / Cold Outreach
    • Always Be Interviewing
    • Interviewing
  • At your job
    • Brag Doc
    • Glue Work
    • Managing U

# Cover Letter / Cold Outreach

I typically don’t subscribe to cover letters. More often than not no one is going to read them and at best won’t make a difference in the decision and at worst will leave a less than great impression. If a company asks for a cover letter I would either skip if it’s optional or don’t apply. It’s not worth your time investment.

However, if the company is small and / or you can reach out to the hiring manager directly, a short “cover letter” in the guise of an email can be a great way to get noticed and show your interest.

Dear {hiring manager},

I’m excited to submit my application for {job title} at {company}. I found this job posting online and after comparing the responsibilities {x}, {x}, {z} to my own skills and interests {a}, {b}, {c}, I jumped at the chance to apply. I believe my {x} years experience in {industry} have allowed me to develop the critical skills necessary to suscceed in this role. Keep it short. No one wants to read a long email from a stranger and you don’t want to spend a good chunk of your time writing something that will be skimmed.

# Interviewing

Be a Story-Teller

There are four stories you have to be able to tell. These stories will answer almost every behavior-based question. - STAR method - Beginning, middle end

Your most successful project or best accomplishment

Knowing this will answer a great success or when you succeeded. It will also answer a time when you hit a tight deadline or tight budget.

Least successful story or time your failed

It will answer questions about time you dealt with difficulties on a project. {} ##### Story about a difficult stakeholder Might be more than one, supervisor, stakeholder, customer, co-worker.

A passion project you’ve worked on

Something you’ve worked on in the past or something you’d like to work on. What gets you up in the morning? What drives you? Where you want to go with your career?

How to ask about the culture

Just as in a behavioral interview, the company is seeing how you behave through experiences, you can ask specific questions to learn how the company behaves in different situations. For example: - How does the company encourage learning an growth? - Tel me about a mistake someone on the team made and how was it handled? - Tell me about a suggestion an employee made that was implemented by the company / team? - What’s a recent projec that excited you?

How to answer ‘tell me about yourself’

Don’t run them through everything in your resume. They already have your resume and ends up being dry. Instead you want to make sure you hit the following four components: your background, your passion, 3 relevant skills or strengths that you’ve built, and what you hope to achieve or get from your next role. You can also share some personal details if appropriate.

Hi I’m {name}, I’ve been a {role} for the past {X} years and I have a background in {domain} and {domain}. One of the things I really love about this job path is {what have you loved about your career path}. The last few roles I’ve had have really helped me hone my skills in {skill 1}, {skill 2}, {skill 3}. And I’m here today because I really want to take my skillset and apply it in the {company domain}. Outside of work I love {hobby 1} and am really into {hobby 2}.

Hi I’m Robert, I’ve been a Data Scientist & Machine Learning Engineer for the past 5 years and have a background in Machine Learning, Distributed Systems, and Geospatial. One of the things I really love about this job path is that there’s never a dull day; there’s always something new to explore in the data or some pipeline to get stuck into and optimize. The last few roles have helped me develop my skills in stakeholder management, technical writing, and cross-team collaboration. I’ve excelled at smaller startups and have developed a wide range of skillsets and I’m here today because I want to go deep in Machine Learning and training on big datasets. Outside of work I love to go climbing and backpacking and am really into cooking at home.

]]> career advice https://bear-toes.pages.dev/posts/career_tips/ Fri, 31 May 2024 23:00:00 GMT Deep Learning Notes https://bear-toes.pages.dev/posts/deep-learning-notes/

# Intuition about deep representation

  • shallower nueral networks need exponentially more () more hidden units vs. a “small” L-layer deep neural network to compute the same function.

Forward and Backward Functions

Given the following generic deep neural network,

Generic Deep Neural Network

A Generic Deep Neural Network

For each layer : ,

Forward Pass:

  • Input:
  • Output: , where:
    • (cache for backward pass)

Backward Pass:

  • Input:
  • Output: , , , where:

Feed Forward Neural Networks Derivations

Parameters & Hyperparameters

Parameters:

Hyperparameters: the parameters that control the Parameters and

  • - learning rate
  • learning rate decay
  • num iterations
  • - num hidden layers
  • - num hidden units
  • - activation functions
  • Momentum
  • - minibatch size
  • regularization

How to learn hyperparamters?

In deep learning, there can be many hyperparameters. With a small amount of hyperparameters, a grid search can be conducted. In deep learning, there can be many hyperparameters, and a grid search becomes untenable. Instead, sampling at random, more unique values are able to be searched.

When searching hyperparamters, a common technique is to first conduct a coarse search over a larger range of values. Then a fine search can be conducted on a smaller region.

Regularization

Bias vs. Variance Tradeoff

Intro to regularization. - Why regularization - What is regularization?

L1 Regularization

{Visualization} {Pytorch Example}

Summary

L2 Regularization

{Visualization} {Pytorch Example}

Summary

Elastic Net Regularization

{Visualization} {Pytorch Example}

Summary

Dropout

Dropout is where, during training, some % of neurons in each layer are zeroed out or “dropped”. The neurons dropped change each batch. Dropout prevents units from co-adapting too much to the data and acts as a sampling strategy since we drop a different set of neurons each time. It effectively forces the net to learn the data without cheating.

{Visualization} {Pytorch Example}

Summary

  • only used during training

Batch Normalization

{Visualization} {Pytorch Example}

Summary - check notes from NN Zero to Heroon this

Other?

Activation Functions

Sigmoid

ReLU

Argmax

Softmax

Softmax regression generalizes logistic regression to classes. If , the softmax reduces to logisitc regression. Softmax is named from the contrast to “Hardmax” or Argmax function.

  • used for multiclass classification
  • normalizes outputs to sum to 1
  • output can be interpreted as probabilities

Loss Function

Where: - is the true label of the class - the summation goes away because all other class outputs are

Cost Function

Optimizers

Gradient Descent

Stochastic Gradient Descent

ADAM

]]> deep learning https://bear-toes.pages.dev/posts/deep-learning-notes/ Sun, 31 Mar 2024 23:00:00 GMT Pub Quiz: Summing Two Random Variables https://bear-toes.pages.dev/posts/summing_random_variables/

# Pub Quiz

Last week, my company held a pub quiz. This was no ordinary pub quiz, it was designed in a way for teams to learn about other teams. One person from each team submitted four questions about their team. We ended up wtih about 60 questions from all across the company. And when I say “all” I mean all. We had questions about Data Science, Commercial, Operations, Engineering, Product, Marketing, Leadership, Human Resources, you name it. As someone who is fairly familiar with all aspects of our products and engineering teams (but by no means an expert in all these areas) by having been here over four years and written my share of documentation, I came into the pub quiz with confidence. Much of that confidence due to expecting to not need to know British pop culture back through the 70s as is standard in an actual pub quiz.

That confidence was quickly smashed with questions like:

  • What was the 2023 Coca Cola Christmas commercial slogan?
  • What is the collective age of the Customer Support team?
  • What month was repo XXX created?

The Data Science round was the trickiest of them all, in no small part because reading the question took up almost all of the response time! After the quiz, my teammate who wrote the Data Science questions shared one that didn’t quite make the cut.

I share the question with the reader here, with only some minor rewording. My answer follows.

# Question

## The preamble

You fit the following univariate regression model using Ordinary Least Squares (OLS)

In the original question, the OLS abbreviation was given without the unabbreviated form. Not sure it would have made the question any easier.

Where the residual, , is Normally distributed with mean 0 and standard deviation 1, that is:

From OLS, you determine:

The next month, your colleague reruns the experiment and collects the same size dataset. They forgot to check the callibration of the machine used to collect the data and as a result the dataset now has a measurement error, . That is,

The measurement error has mean 2 and standard deviation 2.

The residuals in the new sample dataset are also normally distributed.

## And finally the question

What values does your colleague find when running a regression with OLS? (Try solving with a pen and paper):


# Analytical solution

There’s a few things to immediately note:

  1. The hell this is a pub quiz question!
  2. We assume , , , and are independent random variables. If they weren’t, this would be a wee bit trickier (we would need to have information about the joint distributions, i.e. covariances).
  3. We recognize that the new residuals are going to include the previous residuals as well as the new measurement error, .
  4. Whether it was a typo or meant to be tricky, the question states the standard deviation rather than the variance for the Normal distribution. Typically, the notation is uses the variance, . Rewriting into the standard notation:

I’ve highly probably in all likelihood had the same question on a Probability exam during my Statistics postgrad.


The problem boils down to realizing that we have a summation of two independent Normally distributed random variables, and . How do independent Normally distributed random variables sum? If you don’t remember, don’t worry, neither did I. After some revision we know that (Ross 2010):

The assumption that the random variables are independent means the covariance is zero, . We can simply add the means and variances (Ross 2010, 256–57) giving:

where,

Plugging in the values for and , we find the following solution:


Remember, we were asked for not :

We can arguably simplify this a bit further by noting that the bias, , and the mean of the residuals, , are both constants that shift the intercept and can be grouped. Subtracting the mean of the error from the bias…

And rewritting using the standard deviation rather than the variance (like in the original question) we arrive at:

If instead we took 2 as the var in :


# Computational verification

Let’s check our analytical solution through computational methods. First let’s check that we can indeed simply add the means and variances of independent random variables.

Code 
import altair as alt
import pandas as pd
import torch
from sklearn.neighbors import KernelDensity

n = 100_000

err = torch.normal(0, 1, size=(n,))
u = torch.normal(2, 2, size=(n,))

err_computational = err + u
err_analytical = torch.normal(2, 5 ** (0.5), size=(n,))

# torch.histogram returns the endpoints of each bin, we want the middle value
# so we find the mean of each bin range
y, x = torch.histogram(err, density=True)
x = torch.tensor([x[i : i + 2].mean() for i in range(0, len(x) - 1)])
df0 = pd.DataFrame(
    {
        "x": x,
        "density": y,
        "series": ["err"] * len(x),
    }
)

y, x = torch.histogram(u, density=True)
x = torch.tensor([x[i : i + 2].mean() for i in range(0, len(x) - 1)])
df1 = pd.DataFrame(
    {
        "x": x,
        "density": y,
        "series": ["u"] * len(x),
    }
)

y, x = torch.histogram(err_computational, density=True)
x = torch.tensor([x[i : i + 2].mean() for i in range(0, len(x) - 1)])
df2 = pd.DataFrame(
    {
        "x": x,
        "density": y,
        "series": ["err* (computational)"] * len(x),
    }
)

y, x = torch.histogram(err_analytical, density=True)
x = torch.tensor([x[i : i + 2].mean() for i in range(0, len(x) - 1)])
df3 = pd.DataFrame(
    {
        "x": x,
        "density": y,
        "series": ["err* (analytical)"] * len(x),
    }
)
df = pd.concat([df0, df1, df2, df3])

# Plot
colorscheme = [
    "#368BC1",  # blue
    "#F2BB18",  # yellow
    "#BB4430",  # red
    "#8A9A67",  # green
    "#CC771F",  # orange
    "#8B5260",  # purple
]
selection = alt.selection_point(fields=["series"], bind="legend")

chart = (
    alt.Chart(df)
    .mark_line()
    .encode(
        x="x:Q",
        y="density:Q",
        color=alt.Color(
            "series:N", title="Random Variable", scale=alt.Scale(range=colorscheme)
        ).legend(orient="top-right"),
        opacity=alt.condition(selection, alt.value(0.75), alt.value(0.2)),
    )
    .configure(background="#f8f9fa")
    .properties(
        width=650,
        padding=10,
        title={
            "text": f"Summing Two Independent Random Variables (n={n} samples)",
            "subtitle": "err + u = err*",
        },
    )
    .add_params(selection)
)
chart.display()

stats_df = pd.DataFrame(
    {
        "mean": [
            err.mean().item(),
            u.mean().item(),
            err_computational.mean().item(),
            err_analytical.mean().item(),
        ],
        "var": [
            err.var().item(),
            u.var().item(),
            err_computational.var().item(),
            err_analytical.var().item(),
        ],
    },
    index=["ϵ", "u", "ϵ* (computational)", "ϵ* (analytical)"],
)

print(stats_df)
                        mean       var
ϵ                  -0.001720  1.009451
u                   2.005137  3.989215
ϵ* (computational)  2.003417  4.973287
ϵ* (analytical)     2.008872  4.990684

In the chart, we can see is centered on 0 while , , and are all centered on 2. Summing the means worked!

The variance is a bit more difficult to verify from the chart. We can see that and have the same distribution, slightly wider than . The tableThe error between the means and variances is small. If we increase the sample size these errors move closer towards 0.

Now that we’re confident in random variable arithmetic, let’s check the full analytical solution to a comutational one. Below we plot the first measured sample, , with the original signal, , as a backdrop for reference:

Code 
import altair as alt
import pandas as pd
import torch

alt.data_transformers.disable_max_rows()

x = torch.arange(0, 10, 0.01)

# Original Signal
a = 3
b = 2
signal = a + (b * x)

# y Our measured sample
err_mean = 0
err_std = 1
err = torch.normal(err_mean, err_std**2, size=(len(x),))
y = a + (b * x) + err

# y* Analytical
a = 5
b = 2
err_mean = 0
err_std = 5 ** (0.5)
err_analytical = torch.normal(err_mean, err_std, size=(len(x),))
y_analytical = a + (b * x) + err_analytical

# y* Computational
a = 3
b = 2
err_mean = 2
err_std = 5 ** (0.5)
err_computational = torch.normal(err_mean, err_std, size=(len(x),))
y_computational = a + (b * x) + (err + err_computational)

# Data munging
df0 = pd.DataFrame({"x": x, "y": signal, "dataset": ["signal (original)"] * len(x)})
df1 = pd.DataFrame({"x": x, "y": y, "dataset": ["y"] * len(x)})
df2 = pd.DataFrame(
    {"x": x, "y": y_computational, "dataset": ["y* (computational)"] * len(x)}
)
df3 = pd.DataFrame({"x": x, "y": y_analytical, "dataset": ["y* (analytical)"] * len(x)})

df = pd.concat([df0, df1, df2, df3])

# Plot
colorscheme = [
    "#368BC1",  # blue
    "#F2BB18",  # yellow
    "#BB4430",  # red
    "#8A9A67",  # green
    "#CC771F",  # orange
    "#8B5260",  # purple
]
selection = alt.selection_point(fields=["dataset"], bind="legend")
chart = (
    alt.Chart(
        df,
        title="Computational simulation of the problem matches the analytical solution",
    )
    .mark_circle(size=10)
    .encode(
        x="x:Q",
        y="y:Q",
        color=alt.Color(
            "dataset:N", title=None, scale=alt.Scale(range=colorscheme)
        ).legend(orient="top-left"),
        opacity=alt.condition(selection, alt.value(0.75), alt.value(0.2)),
    )
    .configure(background="#f8f9fa")
    .properties(width=650, padding=10)
    .add_params(selection)
)

chart.display()

Note, torch.normal takes the std dev instead of the variance.

To see each series more clearly, click on one in the legend.

The chart shows both the computational analysis verifies the analytical solution while the table above shows the mean and variance of the residuals of the computational anlysis is the same as the analytical solution. The table of statistics for each series also reflects these findings.

# In the end, it didn’t even matter

That took a bit of work. Thankfully, this question didn’t make it into our company pub quiz. Even if it had, I think everyone (sans the author) would have had to ultimately guess the answer. We only had 30 seconds per question. I didn’t even finish reading the problem statement let alone begin tackling it in that time!

In the end, whether this question would have been included or not, wouldn’t have changed the outcome. The quiz was neck and neck between 3 teams for most of the game. My team bouncing between 1 and 3 place. That is until the final round of questions from the Leadership team, which were highly specific to a certain someone.

If we have another pub quiz, I’ll make sure to sling in some equally maladjusted questions. But I’ll post the answers beforehand here.

Watch this space.

Good thing I had the CEO on my team 🥇

References

Ross, Sheldon. 2010. A First Course in Probability. 8th ed. Pearson Prentice Hall.
]]> probability pub quiz https://bear-toes.pages.dev/posts/summing_random_variables/ Sun, 14 Jan 2024 00:00:00 GMT Waltzing with Python’s Walrus Operator https://bear-toes.pages.dev/posts/waltzing_with_the_walrus/ Python 3.8 introduced a new assignment operator with PEP 572 called assignment expressions, a.k.a the walrus operator. The walrus operator uses the new walrus-like syntax :=, to assign variables within an expression.

It’s been out for a few years at this point (at the time of writing Python 3.12 is around the corner) and I’ve found some joy in how it’s helped elegantly shorten some parts of my code. Here are the ways I’ve made use of the walrus operator.

Error Handling

The walrus operator can help reduce repetition and make error handling a bit more streamlined. In the following snippet, func() will return None to represent an error occurred.

x = func()
if not x:
    print("Error message")
    return

Using the walrus operator, the call to func() can be inlined.

if not x := func():
    print("Error message")
    return

Shaving off a single line may seem trivial but those saved lines can add up. For example, when parsing user inputs and performing validation. In the following snippet, we want to validate user inputs x, y, and z. If there’s a validation problem, validate will return a string with a message explaining what is wrong with the input and a message of how to fix. These messages get appended to a list so all validation messages can be printed out together.

validation_errors = []
if msg := validate(x):
    validation_errors += msg

if msg := validate(y):
    validation_errors += msg

if msg := validate(z):
    validation_errors += msg

if validation_errors: # a non-empty list resolves to True
    for error in validation_errors:
        print(error)
    return

Comprehensions

Let’s say we wanted to create a list of results from expensive function call but only results that aren’t None. With a list comprehension, the expensive function would need to be called twice. Not ideal.

y = [
    expensive_function(i) if expensive_function(i) for i in range(0, 10)
]

Of course, you could use normal for loop syntax but it’s a fair bit more verbose, and for illustrative purposes, we’re allergic to verbose.

y = []
for a in range(0, 10):
    x = expensive_function(i)
    if x == 0:
        y.append(x)

The walrus operator plops to the rescue here and allows us to use a list comprehension.

y = [
    x if (x := expensive_function(i)) for i in range(0, 10)
]

This also applies to dictionary comprehensions.

y = {
    i: x if (x := expensive_function(i)) for i in range(0, 10)
}

Do While Loops

A do-while loop was proposed for Python in PEP 315 but was rejected for not providing a material improvement over the following:

while True:
    x = f(a, b) # setup code
    if not x:   # condition
        break
    # loop body using x

A shortened version of do-while loop can be accomplished by having setup code execute once before the loop and moving the condition into a while loop. However, this is error-prone; x = f(a, b) is duplicated for both the setup code and the loop body, and if it needs changing there are now multiple places that must be updated.

x = f(a, b) # setup code
while x:    # condition
    x = f(a, b)
    # loop body using x

With the walrus operator, it can all be inlined to the while condition.

while x := f(a, b): # setup code and condition
    # loop body using x

Pattern Matching

The walrus operator can also be useful in Pattern Matching. Structural Pattern Matching was introduced in Python 3.10 with PEP 622. If you’re not yet familiar see PEP 363 for a tutorial. The walrus can be useful to inline a function call and store the return value in a variable for use in the cases.

match x := f(a, b):
    case 0:
        # do stuff with x
    case 1:
        # do more stuff with x
    case 2:
        # even more doing with x

An Over-the-Top Overuse Example

While the walrus operator is handy for shaving off a few lines of code, inlining too much can make code difficult to reason about. Use it sparingly, especially with other code-golfing operators. For example, with the ternary operator.

height = get_height(name) if (name := get_name(user_id)) else None

I think this can be okay but I also think it’s clearer written long-form,

if name := get_name(user_id):
    height = get_height(name)
else:
    height = None

It could be formatted over multiple lines so it’s just as readible as a normal if / else and to keep the benefits of the ternary usage by only assigning height once, but it’s now a whopping 5 lines.

height = (
    get_height(name)
    if (name := get_name(user_id))
    else None
)

And remember, just because you can doesn’t mean you should write code like below, if you can avoid it. This example is modified from my own code.

params: Dict[str, Dict[str, float]]

sampler_weights = (
    {
        ModeEnum(mode): weight
        for mode, weight in normalize_weights(weights).items()
    }
    if (weights := params.get("sampler", {}).get("weights"))
    else {"x": 0.5, "y": 0.5}
)

Here, I’ve slapped a dictionary comprehension, a ternary operator, and a walrus operator into the same expression. There’s a lot going on, but it’s formatted over multiple lines to help delineate what’s happening. The variable, params, holds the contents of a configuration .toml that I needed to parse some weights from and convert into an dictionary of {enum: weight}.

]]> python https://bear-toes.pages.dev/posts/waltzing_with_the_walrus/ Sat, 12 Aug 2023 23:00:00 GMT