使Keras / Tensorflow输出OneHotCategorical，但是操作没有“梯度”

Question

问题描述

我有作为指标变量的输入x和输出y，其中每一行都是一个随机的一热点向量，它取决于x的值（数据示例如下所示） ）。

我想训练一个模型，该模型本质上以每列权重的形式学习x和y之间的概率关系。模型必须“选择”一个（也只有一个）指示器以输出。我目前的方法是对分类随机变量进行抽样，并产生一个热向量作为预测。

问题是，当我尝试训练Keras模型时出现错误ValueError: An operation has `None` for gradient。

我发现这个错误很奇怪，因为我已经使用Keras和Tensorflow训练了混合网络，并使用tf.contrib.distributions.Categorical，并且没有遇到任何与梯度相关的问题。

代码

实验

import tensorflow as tf
import tensorflow.contrib.distributions as tfd
import numpy as np
from keras import backend as K
from keras.layers import Layer
from keras.models import Sequential
from keras.utils import to_categorical


def make_xy_prob(rng, size=10000):
    rng = np.random.RandomState(rng) if isinstance(rng, int) else rng
    cols = 3
    weights = np.array([[1, 2, 3]])

    # generate data and drop zeros for now
    x = rng.choice(2, (size, cols))
    is_zeros = x.sum(axis=1) == 0
    x = x[~is_zeros]

    # use weights to create probabilities for determining y
    weighted_x = x * weights
    prob_x = weighted_x / weighted_x.sum(axis=1, keepdims=True)
    y = np.row_stack([to_categorical(rng.choice(cols, p=p), cols) for p in prob_x])

    # add zeros back and shuffle
    zeros = np.zeros(((size - len(x), cols)))
    x = np.row_stack([x, zeros])
    y = np.row_stack([y, zeros])
    shuffle_idx = rng.permutation(size)
    x = x[shuffle_idx]
    y = y[shuffle_idx]
    return x, y


class OneHotGate(Layer):
    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel', shape=(1, input_shape[1]), initializer='ones')

    def call(self, x):
        zero_cond = x < 1
        x_shape = tf.shape(x)

        # weight indicators so that more probability is assigned to more likely columns
        weighted_x = x * self.kernel

        # fill zeros with -inf so that zero probability is assigned to that column
        ninf_fill = tf.fill(x_shape, -np.inf)
        masked_x = tf.where(zero_cond, ninf_fill, weighted_x)
        onehot_gate = tf.squeeze(tfd.OneHotCategorical(logits=masked_x, dtype=x.dtype).sample(1))

        # fill gate with zeros where input was originally zero
        zeros_fill = tf.fill(x_shape, 0.0)
        masked_gate = tf.where(zero_cond, zeros_fill, onehot_gate)
        return masked_gate


def experiment(epochs=10):
    K.clear_session()
    rng = np.random.RandomState(2)

    X, y = make_xy_prob(rng)
    input_shape = (X.shape[1], )

    model = Sequential()
    gate_layer = OneHotGate(input_shape=input_shape)
    model.add(gate_layer)
    model.compile('adam', 'categorical_crossentropy')
    model.fit(X, y, 64, epochs, verbose=1)

数据样本

>>> x 
array([[1., 1., 1.],
       [0., 1., 0.],
       [1., 0., 1.],
       ...,
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 0.]])

>>> y
array([[0., 0., 1.],
       [0., 1., 0.],
       [1., 0., 0.],
       ...,
       [0., 0., 1.],
       [1., 0., 0.],
       [1., 0., 0.]])

错误

ValueError: An operation has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.

Answer 1

问题在于OneHotCategorical中执行不连续采样-导致梯度计算失败的事实。为了用连续的（ relaxed ）版本代替这种不连续的采样，可以尝试使用RelaxedOneHotCategorical（基于有趣的Gumbel Softmax技术）。