Tensorflow Implementation of Focal Frequency Loss about focal-frequency-loss HOT 4 CLOSED

from focal_frequency_loss import FocalFrequencyLoss as FFL
from tf_focal_frequency_loss import FocalFrequencyLoss as TFFL

# Pytorch
import torch

ffl = FFL(loss_weight=1.0, alpha=1.0)  # initialize nn.Module class

fake = torch.randn(4, 3, 64, 64)
real = torch.randn(4, 3, 64, 64)

pt_loss = ffl(fake, real)  # calculate focal frequency loss


# Tensorflow
import tensorflow as tf

tffl = TFFL(loss_weight=1.0, alpha=1.0)  # initialize tf.keras.layers.Layer class

fake = tf.convert_to_tensor(fake.numpy(), tf.float32)
real = tf.convert_to_tensor(real.numpy(), tf.float32)

tf_loss = tffl(fake, real)  # calculate focal frequency loss

# Check
import numpy as np
assert np.isclose(pt_loss.numpy(), tf_loss.numpy()) , "Results don't match"