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)
