| import tensorflow.compat.v2 as tf | |
| from absl import flags | |
| FLAGS = flags.FLAGS | |
| class Whitening1D(tf.keras.layers.Layer): | |
| def __init__(self, eps=0, **kwargs): | |
| super(Whitening1D, self).__init__(**kwargs) | |
| self.eps = eps | |
| def call(self, x): | |
| bs, c = x.shape | |
| x_t = tf.transpose(x, (1, 0)) | |
| m = tf.reduce_mean(x_t, axis=1, keepdims=True) | |
| f = x_t - m | |
| ff_apr = tf.matmul(f, f, transpose_b=True) / (tf.cast(bs, tf.float32) - 1.0) | |
| ff_apr_shrinked = (1 - self.eps) * ff_apr + tf.eye(c) * self.eps | |
| sqrt = tf.linalg.cholesky(ff_apr_shrinked) | |
| inv_sqrt = tf.linalg.triangular_solve(sqrt, tf.eye(c)) | |
| f_hat = tf.matmul(inv_sqrt, f) | |
| decorelated = tf.transpose(f_hat, (1, 0)) | |
| return decorelated | |
| def w_mse_loss(x): | |
| """ input x shape = (batch size * num_samples, proj_out_dim) """ | |
| w = Whitening1D() | |
| num_samples = FLAGS.num_samples | |
| num_slice = num_samples * FLAGS.train_batch_size // (2 * FLAGS.proj_out_dim) | |
| x_split = tf.split(x, num_slice, 0) | |
| for i in range(num_slice): | |
| x_split[i] = w(x_split[i]) | |
| x = tf.concat(x_split, 0) | |
| x = tf.math.l2_normalize(x, -1) | |
| x_split = tf.split(x, num_samples, 0) | |
| loss = 0 | |
| for i in range(num_samples - 1): | |
| for j in range(i + 1, num_samples): | |
| v = x_split[i] * x_split[j] | |
| loss += 2 - 2 * tf.reduce_mean(tf.reduce_sum(v, -1)) | |
| loss /= num_samples * (num_samples - 1) // 2 | |
| return loss | |