TensorFlow中对三维数据使用决定系数R2score

问题描述：Tensorflow中提供的R2Score只能接受2维数据（batch， feature），如何对一个三维数据实现R2score的计算？tensorflow传递给模型训练的输入形式什么时候第一个维度是确定的batch值，而不是显示为None？

这是tensorflow中R2Score源码中的一部分，这里判定的模型的输入形状必须是二维且最后一个维度不能为None

本人在搭建LSTM模型时，输入维度为（batch, timestep, feature）,想计算在每个timestep维度上的决定系数，以（2， 200， 5）为例，我相对每条长为200的序列计算预测值与实际值的决定系数，然后再特征维度与batch维度累积求和之后再平均，即2*5=10个决定系数的平均。

自定义了一个 R_squared 类，代码如下：

class R_squared(tf.keras.metrics.Metric):
    def __init__(self, name='r_squared', **kwargs):
        super(R_squared, self).__init__(name=name, **kwargs)
        self.total_samples = self.add_weight(name='total_samples', initializer='zeros')
        self.total_r_squareds = self.add_weight(name='total_r_squareds', initializer='zeros')
        

    def update_state(self, y_true, y_pred, sample_weight=None):
        total_samples = tf.cast(y_true.shape[0] * y_true.shape[2], tf.float32)
        self.total_samples.assign_add(total_samples)

        for i in range(y_true.shape[0]):
            for j in range(y_true.shape[2]):
                true = tf.cast(y_true[i, :, j], tf.float32)
                pred = tf.cast(y_pred[i, :, j], tf.float32)
        
                # Calculate sum of squares of y_true
                sum_of_squares = tf.reduce_sum(tf.square(true - tf.reduce_mean(true)))

                # Calculate sum of residuals
                # self.sum_of_residuals.assign_add(tf.reduce_sum(tf.square(true - pred)))
                sum_of_residuals = tf.reduce_sum(tf.square(true - pred))
                score = 1 - (sum_of_residuals / sum_of_squares)
                self.total_r_squareds.assign_add(score)

    def result(self):
        return (self.total_r_squareds / self.total_samples)

    def reset_state(self):
        # Reset the internal state of the metric
        self.total_samples.assign(0)
        self.total_r_squareds.assign(0)

切记，使用这个方法有一个必须满足的要求：

传递给模型的输入形式必须是显示的例如（2， 200， 5），而不能够是（None, 200, 5）,因此，当model.fit时，传递的batch_size必须是训练集和验证集样本数的公约数，即能被batch_size整除，

否则模型会受到动态的batch，从而显示为None。这个问题卡了我一天半才发现，真的是，绝！

求解最大公约数的代码如下：

import math

batch = math.gcd(a, b)