在Keras中使用CRF进行二值分割时出现形状不兼容错误

Alright, let's figure out why you're hitting that shape mismatch error and fix your CRF integration for binary segmentation.

Why the Error Happens #

The CRF layer from keras_contrib is designed for sequence tasks (like NLP) by default—it expects inputs in the shape (batch_size, sequence_length, num_classes). But your decoder outputs a 2D spatial tensor (batch_size, 224, 336, 1) for segmentation. When you feed this into the standard CRF layer, it misinterprets the spatial dimensions, dropping the width dimension entirely and leading to that shape mismatch.

Correct Ways to Add CRF for Image Segmentation #

Here are two reliable approaches to integrate CRF into your segmentation network properly:

1. Reshape Input to Match Sequence CRF Expectations (Quick Fix)

If you want to stick with the keras_contrib CRF, you can reshape your decoder output to flatten the spatial dimensions into a sequence, pass it through the CRF, then reshape back to the original spatial shape. This works, though it treats each pixel as an independent sequence element (losing explicit spatial neighborhood context, but the CRF will still model class transitions).

Modify your code like this:

import keras
from keras.layers import UpSampling2D, Conv2D, Activation, MaxPooling2D, Reshape
from keras.models import Model
from keras_contrib.layers import CRF

img_width, img_height = 336, 224
kernel_size = 7

input = keras.engine.topology.Input(shape=(img_height, img_width, 3))

# Encoder (unchanged)
e = Conv2D(32,(kernel_size,kernel_size),padding='same')(input)
e1 = Activation('relu')(e)
e = MaxPooling2D(pool_size=(2, 2))(e1)
e = Conv2D(64,(kernel_size,kernel_size),padding='same')(e)
e2 = Activation('relu')(e)
e = MaxPooling2D(pool_size=(2, 2))(e2)

# Decoder (unchanged until final layers)
d = UpSampling2D()(e)
d = Conv2D(64,(kernel_size,kernel_size),padding='same')(d)
d = Activation('relu')(d)
d = UpSampling2D()(d)
d = Conv2D(32,(kernel_size,kernel_size),padding='same')(d)
d = Activation('relu')(d)

# For binary segmentation, output 2 channels (background + foreground)
# CRF requires explicit class channels to model transitions
d = Conv2D(2,(1,1),padding='valid')(d)

# Reshape spatial dimensions to sequence: (batch, height*width, 2)
d = Reshape((img_height * img_width, 2))(d)

# Apply CRF
crf = CRF(2, sparse_target=True)
out = crf(d)

# Reshape back to original spatial shape: (batch, height, width, 2)
out = Reshape((img_height, img_width, 2))(out)

# Extract foreground channel and apply sigmoid for final binary output
out = Activation('sigmoid')(out[..., 1:])

autoencoder = Model(inputs=input, outputs=out)

Key notes:

• We switched the final conv layer to output 2 channels—binary segmentation needs both background and foreground classes for the CRF to model class transitions.
• The reshaping step converts the 2D image tensor into a sequence of pixels, which matches the CRF's expected input format.

2. Use a 2D CRF Layer (Better for Spatial Segmentation)

The sequence-based CRF doesn't explicitly model 2D spatial neighborhood relationships. For better segmentation results, use a dedicated 2D CRF layer designed for images. Here's a differentiable implementation you can add to your code for end-to-end training:

import tensorflow as tf
from keras.layers import Layer

class CRF2D(Layer):
    def __init__(self, num_classes, kernel_size=3, **kwargs):
        self.num_classes = num_classes
        self.kernel_size = kernel_size
        super(CRF2D, self).__init__(**kwargs)

    def build(self, input_shape):
        # Trainable kernel to model spatial neighbor dependencies
        self.kernel = self.add_weight(
            name='crf_kernel',
            shape=(self.kernel_size, self.kernel_size, self.num_classes, self.num_classes),
            initializer='glorot_uniform',
            trainable=True
        )
        super(CRF2D, self).build(input_shape)

    def call(self, x):
        # Combine original logits with spatial neighbor dependencies via convolution
        spatial_logits = tf.nn.conv2d(x, self.kernel, strides=[1,1,1,1], padding='SAME')
        return x + spatial_logits

    def compute_output_shape(self, input_shape):
        return input_shape

Then integrate it into your model like this:

# ... (encoder and decoder layers unchanged)

# Output 2 channels for binary segmentation
d = Conv2D(2,(1,1),padding='valid')(d)
# Apply 2D CRF to model spatial relationships
d = CRF2D(2, kernel_size=3)(d)
# Extract foreground channel and apply sigmoid
out = Activation('sigmoid')(d[..., 1:])

autoencoder = Model(inputs=input, outputs=out)

This 2D CRF uses convolution to explicitly model dependencies between adjacent pixels, which is far more appropriate for image segmentation than the sequence-based variant.

Additional Tips #

• If using sparse_target=True, ensure your labels are integer-encoded (0 for background, 1 for foreground) instead of one-hot.
• If end-to-end training isn't a requirement, you can also apply a traditional CRF as post-processing (using libraries like pydensecrf) on your model's output masks. This is often simpler and works well for many segmentation tasks.

内容的提问来源于stack exchange，提问作者Mark