python - 将函数应用于 3D、2 channel Keras 张量的每个元素

Question

我有一个模型，可以处理 2 个相同大小的输入 3D 图像，A和B ，用于更经典的函数，以尝试提高该函数的性能。为了正确训练模型，我需要将该函数应用于每次运行的结果。该函数本身有 2 个值，分别对应 A 中的值。和B在同一坐标p 。该结果将存储在 3D 图像中， C ，与 A 大小相同和B在点p 。其经典实现将执行 for循环遍历所有坐标并对每一对应用该函数。不幸的是，这种方法不适用于训练 Keras 模型，因为函数的输出必须反馈到前一层的权重。

Input -(A, B)-> Model -(A', B')-> Function(A'[p], B'[p]) -(C[p])-> Result

我尝试为此编写一个自定义 Keras 层。该层接受 4D 张量 (channel, z, y, x)并应返回形状为 (1, z, y, x) 的张量.

目前看起来像:

# imports

def function(x: [float, float]) -> float:
    # a -> x[0], b -> x[1]
    # Calculate
    return c

class CustomLayer(Layer):
    # ... __init__ and build
    def call(self, inputs, **kwargs):
        # All samples, channel n ([::][n])
        # We stack the tensors in such a way because map_fn() maps the top most axis to the function.
        # This way a tensor of shape (n_voxels, 2) is created and the values are delivered in pairs to the function
        map_input = K.stack([K.flatten(inputs[::][0]), K.flatten(inputs[::][1]), axis=1])
        result = K.map_fn(lambda x: function(x), map_input)
        result = K.reshape(result, K.constant([-1, 1, inputs.shape[2], inputs.shape[3], inputs.shape[4]], dtype=tf.int32))
        return result

不幸的是，这种方法严重减慢了训练速度。最后没有自定义层的模型每个周期需要大约 45 分钟的训练时间，而带有自定义层的模型则需要大约 120 小时每个纪元。

模型和函数本身可以执行所需的任务，但会出现重大错误，但是我想看看是否可以将两者结合起来以获得更好的结果。

<小时/>在我的例子中，现实生活中的用途是对 Dual Energy CT中的 Material 进行分解。 。您可以通过假设 3 种已知 Material 来计算每个体素的 Material 分数 mat1 , mat2 , mat3和未知样本 sample .

通过一些计算，您可以将这个未知样本分解为每种已知 Material 的分数 f1 , f2 , f3 ( f1 + f2 + f3 == 1.0)。我们只对f3感兴趣，因此省略了其他分数的计算。

实际代码:

"""
DECT Decomposition Layer
"""

import numpy as np
import tensorflow as tf

from keras import backend as K
from keras.layers import Layer
from keras.constraints import MinMaxNorm


def _intersect(line1_start, line1_end, line2_start, line2_end, ):
    """
    Find the intersection point between 2 lines
    """
    # https://en.wikipedia.org/wiki/Line%E2%80%93line_intersection

    # Tensorflow's tensors need a little more lines to unpack
    x1 = line1_start[0]
    y1 = line1_start[1]
    x2 = line1_end[0]
    y2 = line1_end[1]
    x3 = line2_start[0]
    y3 = line2_start[1]
    x4 = line2_end[0]
    y4 = line2_end[1]

    px = (x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)
    px /= (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
    py = (x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)
    py /= (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)

    return K.stack([px, py])


def _decomp(sample, mat1, mat2, mat3):
    """
    Decomposition of a sample into 1 fraction of material 3
    """
    # Calculate the sample lines' ends
    sample3 = sample + (mat2 - mat3)

    # Calculate the intersection points between the sample lines and triangle sides
    intersect3 = _intersect(sample, sample3, mat1, mat2)

    # Find out how far along the sample line the intersection is
    f3 = tf.norm(sample - intersect3) / tf.norm(sample - sample3)

    return f3


class DectDecompoLayer(Layer):
    def __init__(self, mat1, mat2, mat3, **kwargs):
        self.mat1 = K.constant(mat1)
        self.mat2 = K.constant(mat2)
        self.mat3 = K.constant(mat3)

        super(DectDecompoLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        super(DectDecompoLayer, self).build(input_shape)

    def call(self, inputs, **kwargs):
        map_input = K.stack([K.flatten(inputs[::][0]), K.flatten(inputs[::][1])], axis=1)
        result = K.map_fn(lambda x: _decomp(x, self.mat1, self.mat2, self.mat3), map_input)
        result = K.reshape(result, K.constant([-1, 1, inputs.shape[2], inputs.shape[3], inputs.shape[4]], dtype=tf.int32))
        return result

    def compute_output_shape(self, input_shape):
        return input_shape[0], 1, input_shape[2], input_shape[3], input_shape[4]

Answer 1

好的。首先非常奇怪的是你正在采取“样本”，而不是“ channel ”。

命令inputs[::]准确返回inputs，并且inputs[::][0]等于输入[0]。

因此，无论您的批量大小有多大，您都只训练两个样本。

也就是说，您所需要的只是:

• 假设输入形状为(batch, 2, size, size, size)
• 假设 matN 的形状为 (1, 2, size, size, size)，完全或 (batch, 2, size, size, size)

def call(self, inputs, **kwargs): #shape (batch, 2, size, size, size)
    sample3 = inputs + (self.mat2 - self.mat3) 

    #all shapes (batch, size, size, size)
    x1 = inputs[:,0]
    y1 = inputs[:,1]
    x2 = sample3[:,0]
    y2 = sample3[:,1]

    #all shapes (1, size, size, size) or (batch, size, size, size)
    x3 = self.mat1[:,0]
    y3 = self.mat2[:,1]
    x4 = self.mat2[:,0]
    y4 = self.mat2[:,1]


    #all shapes (batch, size, size, size)
    px = (x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)
    px /= (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
    py = (x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)
    py /= (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)

    #proceed with the rest

警告，平行线可能被零除。

我推荐某种K.switch，例如:

denominator = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
K.switch(
    K.less(K.abs(denominator), K.epsilon()), 
    denominator + K.sign(denominator)*K.epsilon(), 
    denominator)
px /= denominator