# A Partial Replication of MaskFormer in TensorFlow on TPUs for the TensorFlow Model Garden

arXiv: https://arxiv.org/abs/2404.18801 · HTML: https://mlsanity.com/paper/2404.18801/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 25 references that are in the map.
- First posted on arXiv 29 Apr 2024.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation](https://mlsanity.com/paper/1511.00561/) (2015)
- [Per-Pixel Classification is Not All You Need for Semantic Segmentation](https://mlsanity.com/paper/2107.06278/) (2021)
- [Panoptic Segmentation](https://mlsanity.com/paper/1801.00868/) (2018)
- [UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation](https://mlsanity.com/paper/1912.05074/) (2019)
- [Improving Reproducibility in Machine Learning Research (A Report from the NeurIPS 2019 Reproducibility Program)](https://mlsanity.com/paper/2003.12206/) (2020)
- [Pre-Trained Models: Past, Present and Future](https://mlsanity.com/paper/2106.07139/) (2021)
- [An Empirical Study of Pre-Trained Model Reuse in the Hugging Face Deep Learning Model Registry](https://mlsanity.com/paper/2303.02552/) (2023)
- [Challenges and Practices of Deep Learning Model Reengineering: A Case Study on Computer Vision](https://mlsanity.com/paper/2303.07476/) (2023)
- [Reusing Deep Learning Models: Challenges and Directions in Software Engineering](https://mlsanity.com/paper/2404.16688/) (2024)
- [An Empirical Study of Challenges in Converting Deep Learning Models](https://mlsanity.com/paper/2206.14322/) (2022)
- [An Experience Report on Machine Learning Reproducibility: Guidance for Practitioners and TensorFlow Model Garden Contributors](https://mlsanity.com/paper/2107.00821/) (2021)
- [Discrepancies among Pre-trained Deep Neural Networks: A New Threat to Model Zoo Reliability](https://mlsanity.com/paper/2303.02551/) (2023)

## Built on by (most cited first, 1 of 1)

- [Improving the Reproducibility of Deep Learning Software: An Initial Investigation through a Case Study Analysis](https://mlsanity.com/paper/2505.03165/) (2025)

## Abstract (arXiv)

This paper undertakes the task of replicating the MaskFormer model a universal image segmentation model originally developed using the PyTorch framework, within the TensorFlow ecosystem, specifically optimized for execution on Tensor Processing Units (TPUs). Our implementation exploits the modular constructs available within the TensorFlow Model Garden (TFMG), encompassing elements such as the data loader, training orchestrator, and various architectural components, tailored and adapted to meet the specifications of the MaskFormer model. We address key challenges encountered during the replication, non-convergence issues, slow training, adaptation of loss functions, and the integration of TPU-specific functionalities. We verify our reproduced implementation and present qualitative results on the COCO dataset. Although our implementation meets some of the objectives for end-to-end reproducibility, we encountered challenges in replicating the PyTorch version of MaskFormer in TensorFlow. This replication process is not straightforward and requires substantial engineering efforts. Specifically, it necessitates the customization of various components within the TFMG, alongside thorough verification and hyper-parameter tuning. The replication is available at: https://github.com/PurdueDualityLab/tf-maskformer/tree/main/official/projects/maskformer