在TPU上使用TensorFlow对MaskFormer进行部分复现以用于TensorFlow模型花园
计算机视觉与模式识别
2024-04-30 v1 机器学习
软件工程
摘要
本文承担了在TensorFlow生态系统中复现MaskFormer模型的任务，该模型是一种最初使用PyTorch框架开发的通用图像分割模型，并专门针对在张量处理单元（TPUs）上执行进行了优化。我们的实现利用了TensorFlow模型花园（TFMG）中可用的模块化结构，包括数据加载器、训练编排器以及各种架构组件，并针对MaskFormer模型的规格进行了定制和调整。我们解决了复现过程中遇到的关键挑战，包括不收敛问题、训练速度慢、损失函数的调整以及TPU特定功能的集成。我们验证了复现的实现，并在COCO数据集上展示了定性结果。尽管我们的实现达到了端到端可复现性的部分目标，但我们在TensorFlow中复现PyTorch版本的MaskFormer时遇到了挑战。这个复现过程并不简单，需要大量的工程工作。具体来说，它需要定制TFMG中的各种组件，同时进行彻底的验证和超参数调优。该复现代码可在以下地址获取：https://github.com/PurdueDualityLab/tf-maskformer/tree/main/official/projects/maskformer
引用
@article{arxiv.2404.18801,
title = {A Partial Replication of MaskFormer in TensorFlow on TPUs for the TensorFlow Model Garden},
author = {Vishal Purohit and Wenxin Jiang and Akshath R. Ravikiran and James C. Davis},
journal= {arXiv preprint arXiv:2404.18801},
year = {2024}
}