请求开发Keras中non-trainable的2D局部最大值搜索求和层

Got it, let's break down how to build this custom non-trainable Keras layer for your high energy physics use case. This kind of local max + neighbor summation is super common for hit clustering or signal processing in HEP, so let's make sure it's efficient and fits your input shape requirements perfectly.

First, Let's Align on Requirements #

Just to recap so we don't miss anything:

• Input shape: (None, 1, 24, 16) (batch size, single channel, 24px height, 16px width)
• No trainable weights (this layer is purely a data transformation)
• Core logic:
  1. Identify every local maximum in the 2D feature map (a pixel that's larger than all 8 surrounding neighbors)
  2. For each local maximum, sum all pixels in its 3x3 neighborhood (itself + up to 8 adjacent pixels—edge cases will just sum whatever's within bounds)

Implementing the Custom Layer #

We'll use TensorFlow's Keras Layer base class, and leverage built-in TF operations for speed (critical for large HEP datasets). Here's the full code:

import tensorflow as tf
from tensorflow.keras.layers import Layer

class LocalMaxNeighborSum(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # No trainable weights needed, so nothing to initialize here

    def call(self, inputs):
        # Convert input from channel-first (batch, channels, H, W) to channel-last (batch, H, W, channels) for TF ops
        x = tf.transpose(inputs, perm=[0, 2, 3, 1])
        batch_size, height, width, channels = tf.shape(x)

        # Step 1: Detect local maxima in 3x3 neighborhoods
        # Use max pooling to get the highest value in each 3x3 window
        window_max = tf.nn.max_pool2d(x, ksize=3, strides=1, padding='SAME')
        # A pixel is a local max if it equals the max value of its 3x3 window
        local_max_mask = tf.equal(x, window_max)

        # Step 2: Calculate sum of 3x3 neighborhoods for all pixels
        # Avg pool gives sum/9, so multiply by 9 to get the full sum
        # 'SAME' padding handles edge pixels by only summing existing neighbors
        neighbor_sum = tf.nn.avg_pool2d(x, ksize=3, strides=1, padding='SAME') * 9.0

        # Step 3: Keep only sums for local maxima; set others to 0 (adjust if you need a different default)
        result = tf.where(local_max_mask, neighbor_sum, tf.zeros_like(neighbor_sum))

        # Convert back to channel-first shape to match your input format
        result = tf.transpose(result, perm=[0, 3, 1, 2])
        return result

    def compute_output_shape(self, input_shape):
        # Output shape matches input shape since we're doing per-pixel operations
        return input_shape

    def get_config(self):
        # Required for saving/loading your model later
        config = super().get_config()
        return config

Testing the Layer #

Let's quickly verify it works with a sample input:

# Create a test input matching your shape: batch size 1, 1 channel, 24x16
test_input = tf.random.uniform((1, 1, 24, 16), minval=0, maxval=10)

# Initialize and run the layer
layer = LocalMaxNeighborSum()
output = layer(test_input)

# Check that input and output shapes match
print(f"Input shape: {test_input.shape}")
print(f"Output shape: {output.shape}")

# Verify non-max positions have 0 (or your chosen default value)
non_max_pos = tf.where(tf.logical_not(layer(test_input)[0,0,...]))[0]
print(f"Non-max position output value: {output[0,0,non_max_pos[0],non_max_pos[1]]}")

HEP-Specific Notes #

• Edge Handling: The padding='SAME' ensures border pixels are processed correctly—no need to worry about out-of-bounds errors, as the layer will only sum existing neighbors.
• Efficiency: Using TF's built-in pooling ops is way faster than writing custom loops, which is essential for processing large HEP datasets efficiently.
• Customization: If you need a different definition of "local maximum" (e.g., only 4 neighbors instead of 8), you can swap the pooling kernel for a custom mask or adjust the max comparison logic.
• Non-Trainable Guarantee: Since we don't define any trainable weights in the layer, it will stay frozen automatically during model training.

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