frernd_FRN的实现TensorFlow and pytorch

import tensorflow as tfdef FRNLayer(x,tua,beta,gamma,eps=1e-6): """ :param x: input tensor:[BHWC] :param tua: [1,1,1,C] :param beta: [1,1,1,C] :param gamma: [1,1,1,C] :param eps: :return: """ nu2 = tf.reduce_mean(tf.square(x),axis=[1,2], keepdims=True) x = x * tf.rsqrt(nu2 + tf.abs(eps)) return tf.maximum(gamma * x + beta,tua)import torchdef FRN(x,tua,beta,gamma,eps=1e-6): nu2 = torch.mean(torch.square(x),dim=[1,2], keepdim=True) x = x * torch.rsqrt(nu2 + torch.abs(eps)) return torch.maximum(gamma * x + beta,tua)