Problems with PlaidML

[leedrake5]

Just a quick introduction - PlaidML is a backend that can be used to allow GPU-based learning on different hardware (in my case, a Mac with an AMD GPU). If I use PlaidML as a backend:

import os
os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"
from keras import backend as K
from keras import activations, ini

And then attempt to use the the DenseVariational layer:

class DenseVariational(Layer):
    def __init__(self, output_dim, kl_loss_weight, activation=None, **kwargs):
        self.output_dim = output_dim
        self.kl_loss_weight = kl_loss_weight
        self.activation = activations.get(activation)
        super().__init__(**kwargs)
        
    def build(self, input_shape):
        self._trainable_weights.append(prior_params)
        self.kernel_mu = self.add_weight(name='kernel_mu', shape=(input_shape[1], self.output_dim), initializer=initializers.normal(stddev=prior_sigma), trainable=True)
        self.bias_mu = self.add_weight(name='bias_mu', shape=(self.output_dim,), initializer=initializers.normal(stddev=prior_sigma), trainable=True)
        self.kernel_rho = self.add_weight(name='kernel_rho', shape=(input_shape[1], self.output_dim), initializer=initializers.constant(0.0), trainable=True)
        self.bias_rho = self.add_weight(name='bias_rho', shape=(self.output_dim,), initializer=initializers.constant(0.0), trainable=True)
        super().build(input_shape)
        
    def call(self, x):
        kernel_sigma = tf.math.softplus(self.kernel_rho)
        kernel = self.kernel_mu + kernel_sigma * tf.random.normal(self.kernel_mu.shape)
        bias_sigma = tf.math.softplus(self.bias_rho)
        bias = self.bias_mu + bias_sigma * tf.random.normal(self.bias_mu.shape)
        self.add_loss(self.kl_loss(kernel, self.kernel_mu, kernel_sigma) + self.kl_loss(bias, self.bias_mu, bias_sigma))
        return self.activation(K.dot(x, kernel) + bias)
        
    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim)
        
    def kl_loss(self, w, mu, sigma):
        variational_dist = tf.distributions.Normal(mu, sigma)
        return self.kl_loss_weight * K.sum(variational_dist.log_prob(w) - log_mixture_prior_prob(w))

I get the following error when constructing a sequential model:

model.add(DenseVariational(1, kl_loss_weight=kl_loss_weight, activation='sigmoid'))
Traceback (most recent call last):
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/tensor_util.py", line 541, in make_tensor_proto
str_values = [compat.as_bytes(x) for x in proto_values]
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/tensor_util.py", line 541, in
str_values = [compat.as_bytes(x) for x in proto_values]
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/util/compat.py", line 71, in as_bytes
(bytes_or_text,))
TypeError: Expected binary or unicode string, got <tile.Value dense_variational_1/dense_variational_2/dense_variational_3/dense_variational_4/dense_variational_5/dense_variational_6/dense_variational_7/dense_variational_8/dense_variational_9/dense_variational_10/dense_variational_11/dense_variational_12/dense_variational_13/dense_variational_14/dense_variational_15/dense_variational_16/dense_variational_17/dense_variational_18/dense_variational_19/dense_variational_20/dense_variational_21/dense_variational_22/dense_variational_23/dense_variational_24/dense_variational_25/dense_variational_26/dense_variational_27/dense_variational_28/dense_variational_29/kernel_rho Tensor FLOAT32(512, 1)>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "", line 1, in
File "/usr/local/lib/python3.7/site-packages/keras/engine/sequential.py", line 181, in add
output_tensor = layer(self.outputs[0])
File "/usr/local/lib/python3.7/site-packages/keras/engine/base_layer.py", line 457, in call
output = self.call(inputs, **kwargs)
File "", line 17, in call
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_nn_ops.py", line 11532, in softplus
"Softplus", features=features, name=name)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py", line 531, in _apply_op_helper
raise err
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py", line 528, in _apply_op_helper
preferred_dtype=default_dtype)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1297, in internal_convert_to_tensor
ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 286, in _constant_tensor_conversion_function
return constant(v, dtype=dtype, name=name)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 227, in constant
allow_broadcast=True)
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 265, in _constant_impl
allow_broadcast=allow_broadcast))
File "/usr/local/lib/python3.7/site-packages/tensorflow_core/python/framework/tensor_util.py", line 545, in make_tensor_proto
"supported type." % (type(values), values))
TypeError: Failed to convert object of type <class 'plaidml.tile.Value'> to Tensor. Contents: dense_variational_1/dense_variational_2/dense_variational_3/dense_variational_4/dense_variational_5/dense_variational_6/dense_variational_7/dense_variational_8/dense_variational_9/dense_variational_10/dense_variational_11/dense_variational_12/dense_variational_13/dense_variational_14/dense_variational_15/dense_variational_16/dense_variational_17/dense_variational_18/dense_variational_19/dense_variational_20/dense_variational_21/dense_variational_22/dense_variational_23/dense_variational_24/dense_variational_25/dense_variational_26/dense_variational_27/dense_variational_28/dense_variational_29/kernel_rho Tensor FLOAT32(512, 1). Consider casting elements to a supported type.

This may just come with the package if the underlying issue is similar to here, but is there any way to make these layers function without tying them so close to Tensorflow directly?

[krasserm]

Unfortunately, the Keras backend API doesn't provide probability distributions needed here such as tf.distributions.Normal. I could implement them low-level but this wouldn't really scale to more complex distributions. OTOH, tf.random.normal, for example, could be replaced by K.random_normal. If PlaidML has an API for probability distributions, the code should be rather straightforward to migrate I think (but I must admit I have no experience with PlaidML).