def xla_layer_norm_param_grad_job( dy=flow.FixedTensorDef(shape, dtype=dtype), x=flow.FixedTensorDef(shape, dtype=dtype), mean=flow.FixedTensorDef(mean_shape, dtype=dtype), inv_variance=flow.FixedTensorDef(mean_shape, dtype=dtype), ): return flow.layers.layer_norm_param_grad(dy, x, mean, inv_variance, begin_params_axis=params_axis)
def xla_matmul_job( a=flow.FixedTensorDef(a_shape, dtype=dtype), b=flow.FixedTensorDef(b_shape, dtype=dtype), ): out = flow.matmul(a, b, transpose_a=trans_a, transpose_b=trans_b) c = flow.get_variable( "c", shape=out.shape, dtype=flow.float, initializer=flow.ones_initializer(), trainable=True, ) out = flow.math.add_n([out, c]) return out
def trt_avg_pooling_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.nn.avg_pool2d( x, ksize=ksize, strides=strides, padding=padding, data_format=data_format, )
def train_discriminator( z=flow.FixedTensorDef((self.batch_size, 100)), images=flow.FixedTensorDef((self.batch_size, 1, 28, 28)), ): g_out = self.generator(z, trainable=False) g_logits = self.discriminator(g_out, trainable=True) d_loss_fake = flow.nn.sigmoid_cross_entropy_with_logits( flow.zeros_like(g_logits), g_logits, name="Dloss_fake_sigmoid_cross_entropy_with_logits" ) d_logits = self.discriminator(images, trainable=True, reuse=True) d_loss_real = flow.nn.sigmoid_cross_entropy_with_logits( flow.ones_like(d_logits), d_logits, name="Dloss_real_sigmoid_cross_entropy_with_logits" ) d_loss = d_loss_fake + d_loss_real flow.losses.add_loss(d_loss) return d_loss, d_loss_fake, d_loss_real
def train_generator( z=flow.FixedTensorDef((self.batch_size, self.z_dim)), ): g_out = self.generator(z, trainable=True) g_logits = self.discriminator(g_out, trainable=False) g_loss = flow.nn.sigmoid_cross_entropy_with_logits( flow.ones_like(g_logits), g_logits, name="Gloss_sigmoid_cross_entropy_with_logits" ) flow.losses.add_loss(g_loss) return g_loss, g_out
def trt_batch_norm_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): out = flow.layers.batch_normalization(x, axis=axis) c = flow.get_variable( "c", shape=out.shape, dtype=flow.float, initializer=flow.ones_initializer(), trainable=True, ) out = flow.math.add_n([out, c]) return out
def xla_reduce_sum_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
def trt_reshape_like_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), like=flow.FixedTensorDef(like_shape, dtype=dtype), ): return flow.reshape_like(x, like)
def scalar_mul_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.math.multiply(x, scalar)
def trt_add_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return x + y + x
def trt_multiply_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return flow.math.multiply(x, y)
def batch_gather_job( x=flow.FixedTensorDef(input_shape, dtype=dtype), indices=flow.FixedTensorDef(indices_shape, dtype=flow.int32), ): return flow.gather(x, indices, batch_dims=axis)
def xla_gelu_grad_job( x=flow.FixedTensorDef(shape, dtype=dtype), dy=flow.FixedTensorDef(shape, dtype=dtype), ): return flow.math.gelu_grad(x, dy)
def trt_reshape_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.reshape(x, shape)
def trt_bias_add_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), bias=flow.FixedTensorDef(b_shape, dtype=dtype), ): return flow.nn.bias_add(x, bias)
def xla_broadcast_min_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return flow.math.minimum(x, y)
def xla_broadcast_add_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return flow.math.add(x, y)
def xla_gelu_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.math.gelu(x)
def trt_relu_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.math.relu(x)
def trt_matmul_job( a=flow.FixedTensorDef(a_shape, dtype=dtype), b=flow.FixedTensorDef(b_shape, dtype=dtype), ): return flow.matmul(a, b, transpose_a=trans_a, transpose_b=trans_b)
def xla_rsqrt_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.math.rsqrt(x)
def trt_conv2d_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), weight=flow.FixedTensorDef(w_shape, dtype=dtype), ): return flow.nn.conv2d(x, weight, strides, padding, None, data_format, dilation_rate)
def xla_scalar_add_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.math.add(x, scalar)
def trt_sigmoid_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.math.sigmoid(x)
def trt_tanh_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.math.tanh(x)
def trt_identity_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.identity(x)
def trt_reduce_mean_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)): return flow.math.reduce_mean(x, axis=axis, keepdims=keepdims)
def xla_broadcast_mul_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return flow.math.multiply(x, y)
def trt_transpose_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)): return flow.transpose(x, perm=permute)
def xla_broadcast_div_job( x=flow.FixedTensorDef(x_shape, dtype=dtype), y=flow.FixedTensorDef(y_shape, dtype=dtype), ): return flow.math.divide(x, y)