TensorFlow to PyTorch for SLEAP: Is it Worth it?

1 Intro

2 Background

Descriptive information on the two frameworks [1].

Production vs Development graph [6].

Google trends for both frameworks (source).

2.1 TensorFlow Specifications

Training pipeline diagram.

The UNet used in this training pipeline. For a higher resolution, check here. Made with Netron!

A single ".slp" video of 2000 frames of shape (1024, 1024, 1) grayscale
2 tracks/flies with the same 13 keypoints or same skeleton (skeleton is comprised of these 13 keypoints) present for all 2000 frames → 2 flies * 2000 frames = 4000 total training instances

Rotation between -180 and 180 degrees
Normalization (though this isn't done in the pipeline if I'm not mistaken)
Crop based on thorax of the fly
Generate Confidence Map
Batch by 4

standard UNet with 3 up/down blocks
each block consisting of 2 conv, batchnorm, activation function stacks with pooling at the end
1.29M parameters

2.2 The PyTorch Equivalent

2.3 Summary

"same padding" for nn.MaxPool2d and nn.ConvTranspose2d
tf.data.AUTOTUNE [18]

3 Process

3.1 Baseline

Environment
Python==3.10.11
TensorFlow==2.8.4
PyTorch==2.0.0
seed = 42 [19]
Hardware: Tesla T4
Dataset
Rotation(-180, 180)
centroid crop 160x160 on node "thorax"
InstanceConfidenceMapGenerator(sigma=1.5, output_stride=2)
DataLoader (PyTorch)
num_workers = num_cores = 2
pin_memory = True
prefetch_factor = 2
Model
filters = 32
filters_rate = 1.5
down_blocks = 4
stem_blocks = 0
up_blocks = 3
convs_per_block = 2
kernel_size = 3
block_contraction = False
all other arguments in the TensorFlow UNet are default values
Optimizer
Adam(lr=1e-4)
Training
FP16 Mixed Precision (PyTorch)
Batch Size = 4
Epochs = 3
averaged across 5 training runs (though for some I do two averaged 5 training runs)

PyTorch is already faster than TensorFlow
PyTorch's speed is consistent across epochs whereas TensorFlow varies a bit
PyTorch has slightly better performance and convergence time (though insignificant)
Training time for both frameworks were around the same range but the TensorFlow version validation time took much longer
PyTorch UNet reaches an extremely low validation loss after the zeroth epoch

3.2 Improving Upon PyTorch

Optimizing existing PyTorch code [20, 21]: These methods will definitely be used where appropriate as they require little rewriting.
Automated Mixed Precision (AMP)
Max batch size, multiple workers, prefetching, pinned memory
removing bias weights
Avoiding unnecessary GPU-CPU synchronizations
Setting benchmark to True
gradient accumulation and checkpointing
Setting gradients to None instead of 0
DistributedDataParallel for multi-GPU computing
8-bit optimizers [28]
PyTorch Ignite [22]
PyTorch Lightning [23]
PyTorch Fabric [24]
HuggingFace Accelerate [25]
MosaicML Composer [26]
Microsoft DeepSpeed [27]

3.3 Optimizing Existing PyTorch Code

Automated Mixed Precision (AMP)
prefetch_factor = 2
pin_memory = True
num_workers = cores (number of cores)

3.4 PyTorch Ignite

engine: the main module responsible for training, evaluating, and event handling
handlers: the equivalent of callbacks in TensorFlow —methods or classes that execute at a given point during the script to track, log, or compute some information
metrics: a suite of built-in metrics from CV to NLP to traditional ML
distributed: a lightweight wrapper for distributed training like XLA on TPUs
exceptions: for raising exceptions; unused by the user
utils: utility functions; unused by the user



# Engine and helper methods.
from ignite.engine import Engine, Events, create_supervised_trainer, create_supervised_evaluator
# Metrics.
from ignite.metrics import Accuracy, Loss
# Built-in Handlers.
from ignite.handlers import ModelCheckpoint
# Logger and utility.
from ignite.contrib.handlers import TensorboardLogger, global_step_from_engine

# DEFINING TRAINER AND EVALUATORS

trainer = create_supervised_trainer(model, optimizer, criterion, device)

val_metrics = {
    "accuracy": Accuracy(),
    "loss": Loss(criterion)
}

train_evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)
val_evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)

# ADDING EVENT HANDLERS

@trainer.on(Events.EPOCH_COMPLETED)
def log_validation_results(trainer):
	# code to log validation results

def score_function(engine):
    return engine.state.metrics["accuracy"]

model_checkpoint = ModelCheckpoint(
    "checkpoint",
    n_saved=2,
    filename_prefix="best",
    score_function=score_function,
    score_name="accuracy",
    global_step_transform=global_step_from_engine(trainer), # helps fetch the trainer's state
)
  
# Alternative method to adding event handlers.
val_evaluator.add_event_handler(Events.COMPLETED, model_checkpoint, {"model": model})

# TRAINING

trainer.run(train_loader, max_epochs=5)

3.5 PyTorch Lightning

Main components of PyTorch Lightning.

LightningModule: the main module where your model, optimizer, and train/validation/test forward methods are defined along with any other helper methods for your model or forward methods
Trainer: this module abstracts your training code much like Ignite's engine; includes methods for callbacks, loggers with integrations, training configuration parameters, and sanity check flags

accelerators: accelerator classes to specify hardware (CPU, GPU, TPU, IPU, HPU, CUDA)
callbacks: built-in callbacks like ModelCheckpoint and EarlyStopping
cli: command line tools and parsers
core: an assortment of hooks for distributed-aware training, LightningDataModule for standardizing the dataset, HyperparametersMixin for saving hyperparameters (often simply used behind the scenes), and LightningOptimizer (abstracted away from user)
loggers: logger integrations with Comet, Neptune, MLflow, tensorboard, W&B, and plain CSV
plugins: though a bit more abstracted from the user, Lightning has a wide variety of plug-ins
precision: for precision training
environments: modules specific to a particular training environment like Kubeflow, SLURM, etc
io: different checkpointing methods
others: LayerSync and TorchSyncBatchNorm for multiprocessing
profiler: different types of profilers for tracing compute performance
trainer: consists solely of the aforementioned Trainer class
strategies: different strategies for distributed and multi-processing training
tuner: consists of only the Tuner class; equivalent to a hyperparameter tuning class in scikit-learn
utilities: helper and utility methods; see below

Utility module in PyTorch Lightning.

import lightning.pytorch as pl

# define the LightningModule
class LitAutoEncoder(pl.LightningModule):
    def __init__(self, *args, **kwargs):
        super().__init__()
	# define model code here

    def forward(self, *args, **kwargs):
	# define model forward pass here
	return output

    def training_step(self, batch, batch_idx):
        # define training step code here
	# can optionally log
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
	# configure optimizer
        optimizer = optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def validation_step(self, batch, batch_idx):
	# define validation step code here
	return loss

    def test_step(self, batch, batch_idx):
	# define test step code here
	return loss

# init the autoencoder
autoencoder = LitAutoEncoder(encoder, decoder)

trainer = pl.Trainer(limit_train_batches=100, max_epochs=1)
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

PyTorch Lightning seems a bit slower than PyTorch with half precision (FP16).
Both runs are consistent across epochs
It seems that the difference in training time for full precision between Lightning and vanilla PyTorch is negligible
Setting num_workers=0 and prefetch_factor=0, from my experiments, did not accelerate PyTorch Lightning or fix the issue
Lightning is slower than PyTorch by about ~6 seconds

3.6 PyTorch Lightning Fabric









PyTorch Fabric is a bit faster than Lightning, but this most likely has to do with the greater amount of boilerplate PyTorch code
The difference between Lightning and vanilla PyTorch was about 5-6 seconds per epoch whereas Fabric is roughly between 2-3 seconds
Fabric is slower than PyTorch by about ~2 seconds

3.7 HuggingFace Accelerate

Example code of HuggingFace Accelerate.



HuggingFace Accelerate's performance is comparable to that of PyTorch Fabric's
Accelerate is slower than PyTorch by about ~4 seconds

3.8 MosaicML Composer

Cost comparison for training models on AWS.



composer: base class for algorithms, callbacks, ComposerModel (similar to LightningModule, dataset wrappers and tools, and the Trainer, a training loop abstraction similar to the Engine in Ignite and Trainer in PyTorch Lightning including integrations with popular libraries like TIMM and experiment tracking tools like W&B



algorithms: efficient training methods as classes
callbacks: built-in callbacks like model checkpointing
core and core.types: internal functions
datasets: built-in datasets like MNIST
devices:
functional: a module that consists of training optimization methods that are functionally applied to your model; more can be found here 



loggers: built-in integration loggers much like Lightning, Fabric, and Accelerate
loss: collection of custom loss functions
metrics: built-in metrics
models: built-in CV and NLP models
optim: learning rate schedulers
profiler: profiler for run time debugging
utils and utils.dist: internal utility functions

MosaicML Composer (with just squeeze and excite) paired with vanilla PyTorch is only ~2 seconds slower
performance is generalized better

3.9 Microsoft DeepSpeed



Mixed Precision
scalability from one GPU to multiple
CPU Offloading: offload layers that require less GPU compute time onto CPU, prioritizing bottleneck layers on GPU
Pipeline Parallelism: multi-machine setups for model parallelism often have idle time periods, pipeline parallelism takes advantage of this idle wait time by passing additional data through certain machines; refer to the below images

Model parallelism on 4 machines.

Model parallelism on 4 machines with PipeDream (pipeline parallelism).

Megatron-LM: one interesting finding from this paper was that, by leveraging linear algebra, dividing a matrix multiplication between 2 large matrices along the column axis can lend itself to more efficient training; refer to the image below

Dividing a matrix A by column is advantageous for distributed training.

ZeRO: a method that partitions activations, optimizer states, and gradients across parallel data processes instead of copying them across for more efficient memory usage; also leverages Constant Buffer Optimization (CBO) and Contiguous Memory Optimization (CMO)
ZeRO-Offload: offloads optimizer memory and computation from GPU to host CPU
Sparse Attention: methods that zeroes out certain parts of the attention matrices; leads to improved context length and is memory-efficient
1-bit optimizers: compression applied to optimizers
Smart Gradient Accumulation/gradient clipping/checkpointing
Curriculum Learning: train model on easier examples to begin with and gradually tackle harder examples

4 Method



PyTorch Wrappers and their compatibility with experiment tracking applications.

Compatibility graph.



5 Results



6 Discussion

7 Conclusion

8 References