A Partial Replication of MaskFormer in TensorFlow on TPUs for the TensorFlow Model Garden
CoRR(2024)
Abstract
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
MoreTranslated text
AI Read Science
Must-Reading Tree
Example
![](https://originalfileserver.aminer.cn/sys/aminer/pubs/mrt_preview.jpeg)
Generate MRT to find the research sequence of this paper
Chat Paper
Summary is being generated by the instructions you defined