ista_tf.py 文件源码

def _get_step(self, inputs):
        Z, X, lmbd = self.inputs
        K, p = self.D.shape
        L = self.L
        with tf.name_scope("step_ISTA"):
            self.S = tf.constant(np.eye(K, dtype=np.float32) - self.S0/L,
                                 shape=[K, K], name='S')
            self.We = tf.constant(self.D.T / L, shape=[p, K],
                                  dtype=tf.float32, name='We')
            B = tf.matmul(X, self.We, name='B')
            hk = tf.matmul(Z, self.S) + B
            step = soft_thresholding(hk, lmbd / L)
            dz = tf.reduce_mean(tf.reduce_sum(
                tf.squared_difference(step, Z), reduction_indices=[1]))
        return step, dz