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)
Example #2
0
 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
Example #3
0
 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,
     )
Example #4
0
        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
Example #5
0
 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
Example #6
0
 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
Example #7
0
 def xla_reduce_sum_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
Example #8
0
 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)
Example #9
0
 def scalar_mul_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.math.multiply(x, scalar)
Example #10
0
 def trt_add_job(
     x=flow.FixedTensorDef(x_shape, dtype=dtype),
     y=flow.FixedTensorDef(y_shape, dtype=dtype),
 ):
     return x + y + x
Example #11
0
 def trt_multiply_job(
     x=flow.FixedTensorDef(x_shape, dtype=dtype),
     y=flow.FixedTensorDef(y_shape, dtype=dtype),
 ):
     return flow.math.multiply(x, y)
Example #12
0
 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)
Example #13
0
 def xla_gelu_grad_job(
         x=flow.FixedTensorDef(shape, dtype=dtype),
         dy=flow.FixedTensorDef(shape, dtype=dtype),
 ):
     return flow.math.gelu_grad(x, dy)
Example #14
0
 def trt_reshape_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.reshape(x, shape)
Example #15
0
 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)
Example #16
0
 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)
Example #17
0
 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)
Example #18
0
 def xla_gelu_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.math.gelu(x)
Example #19
0
 def trt_relu_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.math.relu(x)
Example #20
0
 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)
Example #21
0
 def xla_rsqrt_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.math.rsqrt(x)
Example #22
0
 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)
Example #23
0
 def xla_scalar_add_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.math.add(x, scalar)
Example #24
0
 def trt_sigmoid_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.math.sigmoid(x)
Example #25
0
 def trt_tanh_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.math.tanh(x)
Example #26
0
 def trt_identity_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.identity(x)
Example #27
0
 def trt_reduce_mean_job(x=flow.FixedTensorDef(x_shape, dtype=dtype)):
     return flow.math.reduce_mean(x, axis=axis, keepdims=keepdims)
Example #28
0
 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)
Example #29
0
 def trt_transpose_job(x=flow.FixedTensorDef(input_shape, dtype=dtype)):
     return flow.transpose(x, perm=permute)
Example #30
0
 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)