TensorFlow 中的多类 U-Net 分割

Question

我已经看到了这个问题的各种变体，但仍在努力正确实施它。我有带有 4 classes（0-背景、1-组织类型 1、2-组织类型 2、3-莫名其妙地跳过和 4-组织类型 4...BrATs 数据集）的真值分割掩码的大脑 MRI 图像)

我实现了一个基本的 U-Net 架构，但我无法将其扩展到非二进制 class化。特别是损失函数。

这是我实现的，但我显然忽略了重要的细节：

[...]
output = tf.layers.conv2d_transpose(
conv18,
filters=5,
kernel_size=1,
strides=1,
padding='same',
data_format='channels_last',
activation=None,
use_bias=True,
kernel_initializer=None,
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=tf.contrib.layers.l2_regularizer(reg),
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
trainable=True,
name='output',
reuse=None
)

我认为 (0,1,2,3,4) 个可能的掩码值的 5 个过滤器是正确的。然后我使用了以下损失函数：

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
_sentinel=None,
labels=label,
logits=output,
name='cross_ent_loss'
)

return tf.reduce_mean(loss)

logits 将通过上面的输出，标签将是我堆叠的蒙版图像 [n_batch、x_dim、y_dim、1]。查看文档，我知道我没有将正确的张量传递给标签。

我的做法是否正确？如何使用 1 个蒙版图像中包含的多个 class 标签实现损失？

Answer 1

我直接从文档中忽略的东西 tf.nn.sparse_softmax_cross_entropy_with_logits:

labels: Tensor of shape [d_0, d_1, ..., d_{r-1}] (where r is rank of labels and result) and dtype int32 or int64. Each entry in labels must be an index in [0, num_classes). Other values will raise an exception when this op is run on CPU, and return NaN for corresponding loss and gradient rows on GPU.

因此将标签更改为形状 [-1]

label_reshape = tf.reshape(label, [-1])

修复了！

TensorFlow 中的多类 U-Net 分割

Multiclass U-Net segmentation in TensorFlow

python

image-segmentation

conv-neural-network

tensorflow