U-Net 多分类实现整理

上周发了个博客说多分类应该用一个像素代表不同值的方式来实现，今天又仔细研习了一下 github 上这个 issue，看了作者的回复，感觉还是应该用多通道作为输入。
于是我把unet传入的n_classes参数改成了2，然后让传入的数据形状和unet输出预测的形状一致，把嗅球和白色区域数据用torch.cat叠起来，最后两个通道输出扩充一个通道保存成3通道RGB的形式

我详细阅读了github上这个issue
对于多分类标签的输入输出规格，我从这个issue里作者的回复
大概是这样理解的：
输入的多个标签使用one-hot-encoded对每个像素编码，输出的预测也是按照one-hot-encoded编码。
具体到嗅球分割是这样：
输入图是灰度图，n_channels=1，预测分类有两类，n_classes=2
对于每个像素，使用one-hot-encoded，每个像素有2个值，对于下面4种情况，分别对应含义如下：
(0, 0) – 此像素是背景
(1, 0) – 此像素是嗅球
(0, 1) – 此像素是白色区域
(1, 1) – 此像素是嗅球和白色部分交集
原图大小是896x896，裁剪中间480x480输入。我修改了mridataloader.py中的对应方法，让本来只返回嗅球标签的函数多返回一个白色区域标签，然后label获取部分的代码改为：
label = torch.cat((batch_data[1], batch_data[2]), dim=1).to(self.device)
把两个标签合并成两层，这样的话label的形状就是(2, 480, 480)，和n_classes=2的Unet网络返回的预测标签形状一样。
只要确定两个label形状一致，就可以使用focal loss来进行前向传播和反向传播。然后我在写这部分代码的时候发现了一个问题，就是我之前代码有个bug导致之前一直用的还是原先的L1Loss……使用FocalLoss的代码没有生效。现在是确定用的是FocalLoss了。训练出相似的结果超参数如下：lr=0.0001，迭代5000次左右的模型最佳。我发现使用FocalLoss收敛的速度明显要比L1Loss快很多……效果上FocalLoss跑出来模型的Dice系数更高一点，表现稍好一些。
然后保存的时候，因为不再是1层的结果了，就没法存灰度图，于是把原图和结果都扩展到3x480x480的形状，把one-hot-encode的标签存成了RGB图，这样的话嗅球对应红色，白色区域对应绿色，效果如下：