这是今年发表在MICCAI2021上的论文，也就是我前几篇文章提到最近一直在做的事情。

实不相瞒，这个工作很大程度上是受到CVPR 2020上DSRL那篇文章的启发。当时导师就表示DSRL这个工作很有意思，idea很简单，文章写得很清晰，可以循着思路做下去。

一开始我也没想明白要怎么做下去，所以只是把数据集换成了医疗图像做实验。结果作者说是要公开代码，但是过了一年多了还是没公开，我只好自己复现。我复现的结果和论文很不一样，在医学图像数据集上DSRL反而会导致分割branch什么都学不出来。

我在组会上汇报了这个结果，简单分析了下觉得是因为DSRL里面的FA module只适用于自然图像里面那种多类别输出的任务，像医学图像这种类别很少的用FA规范下特征反而导致特征整体偏向SR branch的输出了，导致segmentation的性能极差（我实验中结果基本就是0）。老师当时不知可否，只说我如果是自己复现的代码，应该用论文中一样的数据集验证下正确性才行。

老师说的很有道理，于是后来几个月我主要都在尝试复现论文的实验结果。看的时候觉得很清晰，复现起来就发现问题很多，没写具体的地方很多。我当时碰到的主要问题就是FTM不知道到底是怎么做的，放在网络的哪个位置。因为论文里确实也没说清楚，我在issue区问的问题作者也回答的不明不白（有空回issue你们怎么就没空公开代码呢），所以我花了很多时间尝试不同的TFM方案。期间还夹杂了一次MICCAI的Grand Challenge的新馆肺炎比赛，搞了我一个月。这比赛用的MONAI框架做的，我验证集上是排名20%，按理说是有资格参加决赛的，结果决赛开始那天没给我发邮件，我没拿到测试集。过了一天我等不了了发邮件问他们，对面才把数据集发给我。不过测试集上我做了一个非常错误的决定，就是我打算用昨天的新模型做测试。这个新模型没有在验证集上测试（当时验证集通道已经关闭），只不过是加了多尺度训练，我认为理论上会性能好一些，结果测试出来直接拍到50%去了，哈哈。

这期间我一直用另一张卡在测试不同的TFM结构，终于有一天给我做出来了。原来这TFM就是简单的用两个branch的输出作为输入的，真给我整无语了，文章里说是用两个branch的特征层做TFM输入害得我一直在试到底是用的那一层特征做input。

验证了复现的结果正确性后，我就直接把代码开源了。到目前为止DSRL的官方代码还是没公开，我开源的这个repo的issue区甚至有人认为我是作者直接质疑我论文。

复现完之后用到医学数据集上，还是不行。这次老师们都信了，但觉得我做这个工作没有motivation。和几个老师激烈讨论一个多小时之后老师们都觉得如果能用到3D分割上就好了，因为3D分割计算量很大吗，用DSRL这个框架能把计算资源需求降下来。当时我说要不我们试试直接用一个降采样的3D图来一次性生成完整的分割mask吧，这样的话就不用割patch了。陈老师说这个确实可以，大家一直都想做的。像是以前自然图像做2D也会分patch的，但现在算力上来了之后基本都直接用整张图了，除非数据量太少用patch来当增广。不用patch的一大优势就是速度快，一次性出一张图；另外就是一次输入一个完整的图像总归比只输入一部分用来训练要好的，就好像输一张猫或者一张狗的图让分类会比你输进去猫和狗的眼睛要容易很多。

但是想实现这个很难。目前很多实验室都是多张1080Ti/2080Ti组成大显存，这样确实可以通过nn.DataParallel把bs拉上去，但是DataParallel的实现原理是每张卡上都放一个model然后把bs分配给每张卡，所以像3D医学图像这种数据可以你256*256*256的图都没法用bs=1直接放到一张卡上，所以唯一的解决方案就是换单卡就有大显存的设备，但大家都知道大显存的单卡都很贵，而且那段时间虚拟币炒得飞起，想买卡都买不到。

我记得我一开始是在CT上做的，想做128*128*128输入然后256*256*256的输出（正常CT都是512*512*d起步的），结果连这个都放不下。所以当时就只好先做64到128的，结果效果差的要死，根本不如用patch的方法。当时很郁闷，觉得这条路可能根本就走不通的，但是仔细分析之后我发现很可能是因为把512直接降采样到64损失太多高频信息了，我看了几个64的slice简直就是糊成了一坨，什么都分不清了，所以之后我就联系医院那边的人，补习影像学的知识。

最后我选用MRI了，因为现在CT太成熟了，数据基本全都是512*512的了，用11G显存的卡真心不好做。但是MRI还好，MRI虽说很多也都是512的高清图了，但是很多时候也有分辨率200多的。我当时问周围的人有没有做过MRI的，有个学长推荐我看看brats，说是真的好用，规整。我打开看了下，每张图都是240*240*155的，spacing也都是1mm，简直是太完美了。刚好brats也是个著名的肿瘤分割比赛，数据也比较有说服力。

另一个好处是brats里面只有脑的部分，其他的都给主办方清理干净了，所以我直接把周围的黑边切掉就得到了192的最终图。我试了下，做96到192的patch-free用11G卡刚刚好，真是天无绝人之路。

之后就是网络设计的问题了，因为DSRL在这个数据集上可以说是完全用不了。不过我们都觉得这种用SR做辅助任务的思路还是可以的，因为实验发现去掉FA的话加上这个SR确实对性能有微弱提升。所以我之后就在此基础上打算往上增加东西，利用好这两个相关的tasks。我一开始是加了个target enhance，也是一个非常非常naive的想法，就是把两个task的结果element-wise了得到lesion区域的高清图，然后希望SR branch能强化这部分区域的高清还原。加上去之后效果惊人。不过汇报的时候老师们都说你这也太简单了，所以后来我仔细分析了下FA为什么不work，然后针对医学图像类别少的特性自己搞了个空间相似度矩阵出来，强化对分割错误后体现在element-wise结果上的惩罚，加上去了也获得了不错的结果。本来我是说这两个模块相当于替代FA了，跟DSRL相似的只有底层都用了segmentation和SR的multitask learning，而multitask learing实际上大家都在用吗，所以跟DSRL其实已经没啥关系了。但是老师们还是觉得怪怪的，因为要手动降采样高分辨图像，总觉得是一种浪费。其实我也这么觉得，但不降采样实在放不下，所以后来我退而求其次，在HR图像中割了一小块patch出来作为guidance了，然后高分辨特征和全局特征concat起来，这样一来还可以和SR形成一个自监督。加上之后整个性能就已经很不错了，老师挺满意的。

然后寒假的时候我就感觉差不多了，马上也到ddl了，赶紧写了一波。第一版老师们看了说写的这也太烂了，写小说呢？于是亲自下场给我改，给我都整感动了。后来很多学长和朋友给我帮助，最后文章就中了。代码我已经开源了，github搜PFSeg就行。

然后最近我在做vision transformer的轻量化，目前也有了点结果。主要就是针对self attention的两个问题做了改进，一个是计算量大，一个是只有自相关，没有引入外相关。现在模型参数量11.47M，是UNet的1/3，transunet的1/11不到，但性能比两个都强不少，之后如果还有进展再来和大家分享。具体操作在发表之前就不过多介绍了，没有同行评议的话其实我心里也没底，就先不发出来误导大家了。