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

Custom loss with loops in tensorflow

我的数据预处理中有一个函数,它在 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 函数是可能的?

我认为这是与您的 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