Пользовательские потери с петлями в тензорном потоке - PullRequest
Новая
       12

Пользовательские потери с петлями в тензорном потоке

1 голос
/ 01 апреля 2019

У меня есть функция предварительной обработки данных, которая выполняет блочный DCT для трехмерных массивов в режиме YCbCr. 

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)

Есть идеи, как это правильно сделать? Я предполагаю, что лямбда-функции могут быть возможны?

1 Ответ

1 голос
/ 01 апреля 2019

Я думаю, что это эквивалент TensorFlow вашей функции NumPy / SciPy: 

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
...