如何为自定义Tendorflow v2.x Optimizer实施稀疏更新支持？

问题描述

我在TensorFlow v2.x中为自定义优化器实现稀疏矩阵支持更新时遇到了问题

# Dense Matrix Implementation
def _resource_apply_dense(self,grad,handle,apply_state=None):

# Sparse Matrix Implementation
def _resource_apply_sparse(self,indices,apply_state=None):

_resource_apply_dense 正常运行。但是 _resource_apply_sparse 无法正常工作。

下面的代码片段是主要功能，可以正常工作，但结果是错误的。

def _resource_apply_sparse(self,apply_state=None):
        # Method 1
        handle_new = tf.gather(handle,indices) 
        gradients_sum = self.get_slot(handle,"gradients_sum")
        grad_norm_sum = self.get_slot(handle,"grad_norm_sum")
        tilde_w = self.get_slot(handle,"tilde_w")
        lr = self.get_slot(handle,"L")
        reward = self.get_slot(handle,"reward")

        lr_update = tf.maximum(lr,tf.abs(grad))
        gradients_sum_update = gradients_sum + grad
        grad_norm_sum_update = grad_norm_sum + tf.abs(grad)
        reward_update = tf.maximum(reward - grad * tilde_w,0)
        new_w = (
            -gradients_sum_update
            / (
                lr_update
                * (
                    tf.maximum(
                        grad_norm_sum_update + lr_update,self._alpha * lr_update
                    )
                )
            )
            * (reward_update + lr_update)
        )
        var_update = handle_new - tilde_w + new_w
        tilde_w_update = new_w

        gradients_sum_update_op = state_ops.assign(gradients_sum,gradients_sum_update)
        grad_norm_sum_update_op = state_ops.assign(grad_norm_sum,grad_norm_sum_update)
        var_update_op = self._resource_scatter_update(handle,var_update)
        tilde_w_update_op = state_ops.assign(tilde_w,tilde_w_update)
        lr_update_op = state_ops.assign(lr,lr_update)
        reward_update_op = state_ops.assign(reward,reward_update)

        return control_flow_ops.group(
            *[
                gradients_sum_update_op,var_update_op,grad_norm_sum_update_op,tilde_w_update_op,reward_update_op,lr_update_op,]
        )

在下面找到pyTest函数。

def run_sparse_sample(iterations,expected,optimizer):
    var_0 = tf.Variable([1.0,2.0])
    var_1 = tf.Variable([3.0,4.0])

    grad_0 = tf.IndexedSlices(tf.constant([0.1]),tf.constant([0]),tf.constant([2]))
    grad_1 = tf.IndexedSlices(tf.constant([0.04]),tf.constant([1]),tf.constant([2]))

    grads_and_vars = list(zip([grad_0,grad_1],[var_0,var_1]))

    for _ in range(iterations):
        optimizer.apply_gradients(grads_and_vars)

    np.testing.assert_allclose(var_0.read_value(),expected[0],atol=2e-4)
    np.testing.assert_allclose(var_1.read_value(),expected[1],atol=2e-4)


def test_sparse_sample_with_default_alpha():
    run_sparse_sample(
        iterations=5,expected=[
            [0.84528893,1.845289],[2.845289,3.845289]
            ],optimizer=COCOB(),)

如果有人可以帮助我，那将是很棒的。让我知道是否需要更多详细信息。

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）