张量流中带有循环的自定义损失

Posted 2023-03-28

【中文标题】张量流中带有循环的自定义损失

【英文标题】：Custom loss with loops in tensorflow

【发布时间】：2019-08-22 09:55:11

【问题描述】：

我的数据预处理中有一个函数，它在 YCbCr 模式下对 3D numpy 数组执行分块 DCT。

def perform_blockwise_dct(img, ratio):

    imsize = img.shape
    dct_blocks = np.zeros(imsize)

    for i in np.r_[:imsize[0]:8]:
        for j in np.r_[:imsize[1]:8]:

            dct_blocks[i:(i+8),j:(j+8), 0] = dct(dct(img[i:(i+8),j:(j+8), 0].T, norm='ortho').T, norm='ortho')
            dct_blocks[i:(i+8),j:(j+8), 1] = dct(dct(img[i:(i+8),j:(j+8), 1].T, norm='ortho').T, norm='ortho')
            dct_blocks[i:(i+8),j:(j+8), 2] = dct(dct(img[i:(i+8),j:(j+8), 2].T, norm='ortho').T, norm='ortho')

    return dct_blocks

为了能够实现自定义均方误差函数，我想反转这个函数。问题是在实现损失函数时，它是一个张量流张量。有一个逆 DCT 函数可供使用。但是，我不知道如何执行等效的双 for 循环来逐块执行。目前它是在整个图像上完成的，如下所示：

    def mse_custom_loss(a, b)
        y = tf.spectral.idct(a[:,:,0], norm='ortho') 
        cb = tf.spectral.idct(a[:,:,1], norm='ortho')
        cr = tf.spectral.idct(a[:,:,2], norm='ortho') 
        a = K.stack([y, cb, cr], axis=-1)

        y = tf.spectral.idct(b[:,:,0], norm='ortho') 
        cb = tf.spectral.idct(b[:,:,1], norm='ortho')
        cr = tf.spectral.idct(b[:,:,2], norm='ortho') 
        b = K.stack([y, cb, cr], axis=-1)

        return mean_square_error(a, b)

关于如何正确执行的任何想法？我认为 lambda 函数可能是一种可能性？

【问题讨论】：

我对 DCT 或一般的信号处理不是很熟悉，但似乎在 perform_blockwise_dct 中您使用 dct 两次来制作 2D DCT，在 mse_custom_loss 中您正在使用idct 一次，这将为您提供最后一维的一维 IDCT。哪一个是对的？另外，对于损失函数，a 和 b 是单个图像（3D 张量）还是批量图像（4D 张量）？

没错，它应该是每个通道的 2D DCT。例如形状 (128, 128, 3) 的图像。换句话说，在自定义损失函数中，它应该只是 2D 而不是 1D。关于 a 和 b，这些是单个图像（3D 张量），而不是图像批次。

【参考方案1】：

我认为这是一个等效于您的 NumPy/SciPy 函数的 TensorFlow：

import tensorflow as tf

def perform_blockwise_dct_tf(img):
    shape = tf.shape(img)
    x, y, c = shape[0], shape[1], shape[2]
    img_res = tf.reshape(img, [x // 8, 8, y // 8, 8, c])
    img_dct1 = tf.spectral.dct(tf.transpose(img_res, [0, 1, 2, 4, 3]), norm='ortho')
    img_dct2 = tf.spectral.dct(tf.transpose(img_dct1, [0, 2, 4, 3, 1]), norm='ortho')
    out = tf.reshape(tf.transpose(img_dct2, [0, 4, 1, 2, 3]), shape)
    return out

一个小测试：

import numpy as np
from scipy.fftpack import dct

def perform_blockwise_dct(img):
    imsize = img.shape
    dct_blocks = np.zeros(imsize, dtype=img.dtype)
    for i in np.r_[:imsize[0]:8]:
        for j in np.r_[:imsize[1]:8]:
            dct_blocks[i:(i+8), j:(j+8), 0] = dct(dct(img[i:(i+8), j:(j+8), 0].T, norm='ortho').T, norm='ortho')
            dct_blocks[i:(i+8), j:(j+8), 1] = dct(dct(img[i:(i+8), j:(j+8), 1].T, norm='ortho').T, norm='ortho')
            dct_blocks[i:(i+8), j:(j+8), 2] = dct(dct(img[i:(i+8), j:(j+8), 2].T, norm='ortho').T, norm='ortho')
    return dct_blocks

np.random.seed(100)
# DCT in TensorFlow only supports float32
img = np.random.rand(128, 256, 3).astype(np.float32)
out1 = perform_blockwise_dct(img)
with tf.Graph().as_default(), tf.Session() as sess:
    out2 = sess.run(perform_blockwise_dct_tf(img))
# There is a bit of error
print(np.allclose(out1, out2, rtol=1e-5, atol=1e-6))
# True