def FlowNnBnJob(
x_full_precision: oft.Numpy.Placeholder(x.shape),
mean: oft.Numpy.Placeholder(mean.shape),
variance: oft.Numpy.Placeholder(variance.shape),
offset: oft.Numpy.Placeholder(offset.shape),
scale: oft.Numpy.Placeholder(scale.shape),
):
with flow.scope.placement(device_type, "0:0"):
x_full_precision += flow.get_variable(
name="v1",
shape=(1, ),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
)
if data_type == "float16":
x = flow.cast(x_full_precision, flow.float16)
else:
x = x_full_precision
y = flow.nn.batch_normalization(x,
mean,
variance,
offset,
scale,
epsilon,
axis=axis)
y = flow.cast(y, flow.float32)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[0]),
momentum=0).minimize(y)
flow.watch_diff(x_full_precision,
test_global_storage.Setter("x_diff"))
return y
def FlowJob(
value: oft.Numpy.Placeholder(value.shape),
bias: oft.Numpy.Placeholder(bias.shape),
):
with flow.scope.placement(device_type, "0:0"):
value += flow.get_variable(
name="v1",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
bias += flow.get_variable(
name="v2",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
if data_type == "float16":
comp_value = flow.cast(value, dtype=flow.float16)
comp_bias = flow.cast(bias, dtype=flow.float16)
else:
comp_value = value
comp_bias = bias
loss = flow.nn.bias_add(comp_value, comp_bias, *flow_args)
if data_type == "float16":
loss = flow.cast(loss, dtype=flow.float)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[0]),
momentum=0).minimize(loss)
flow.watch_diff(value, test_global_storage.Setter("value_diff"))
flow.watch_diff(bias, test_global_storage.Setter("bias_diff"))
return loss
def test_element_wise_mul_job(
x: oft.Numpy.Placeholder(shape, dtype=flow.float),
y: oft.Numpy.Placeholder(shape, dtype=flow.float),
):
with flow.scope.placement(device, "0:0"):
x += flow.get_variable(
name="vx",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
y += flow.get_variable(
name="vy",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x = flow.cast(x, dtype=flow_type)
y = flow.cast(y, dtype=flow_type)
out = flow.math.multiply(x, y)
out = flow.cast(out, dtype=flow.float)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(y, test_global_storage.Setter("y"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def test_fused_scale_tril_fw_bw_job(
x: oft.Numpy.Placeholder(shape, dtype=flow_type)
):
with flow.scope.placement(device, "0:0"):
x_var = flow.get_variable(
name="xv",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x += flow.cast(x_var, dtype=flow_type)
if type_name == "float16":
out = flow.cast(
flow.math.fused_scale_tril(
flow.cast(x, flow.float16), diagonal, scale=scale
),
flow.float,
)
else:
out = flow.math.fused_scale_tril(x, diagonal, scale=scale)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.0001]), momentum=0
).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def DropoutJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=-1,
maxval=1),
trainable=True,
)
if data_type == "float16":
x = flow.cast(flow.cast(x, flow.float16), dtype)
of_out = flow.cast(
flow.nn.dropout(flow.cast(x, flow.float16),
rate=rate,
name="dropout"),
dtype,
)
else:
of_out = flow.nn.dropout(x, rate=rate, name="dropout")
loss = flow.math.square(of_out)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(of_out, test_global_storage.Setter("out"))
flow.watch_diff(of_out, test_global_storage.Setter("out_diff"))
return loss
def SoftmaxJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=-1.0,
maxval=1.0),
trainable=True,
)
x1 = x
x = flow.identity(x)
if data_type == "float16":
loss = flow.cast(
flow.nn.softmax(flow.cast(x, dtype=flow.float16),
axis=axis),
dtype=flow.float,
)
else:
loss = flow.nn.softmax(x, axis=axis)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
total_loss = loss * x1
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(total_loss)
return loss
def test_masked_fill_fw_bw_job(
x: oft.Numpy.Placeholder(x_shape, dtype=flow_type),
mask: oft.Numpy.Placeholder(mask_shape, dtype=flow_type),
):
with flow.scope.placement(device, "0:0"):
y = flow.get_variable(
name="vx",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x += flow.cast(y, flow_type)
mask = flow.cast(mask, dtype=flow.int8)
if type_name == "float16":
out = flow.cast(
flow.masked_fill(flow.cast(x, flow.float16), mask, value),
flow.float,
)
else:
out = flow.masked_fill(x, mask, value)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(out)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(out, test_global_storage.Setter("out"))
flow.watch_diff(out, test_global_storage.Setter("out_diff"))
return out
def expandJob(
of_input: tp.Numpy.Placeholder(shape=input.shape,
dtype=flow.float32),
multipler: tp.Numpy.Placeholder(shape=gout.shape,
dtype=flow.float32,
batch_axis=diff),
) -> tp.Numpy:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=of_input.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
name="v",
)
input_x = v + of_input
flow.watch_diff(input_x, assert_prediction_grad)
x_fp32 = flow.cast(input_x, flow.float32)
x_fp16 = flow.cast(input_x, dtype=flow.float16)
y_fp16 = flow.expand(x_fp16, expand_dim)
y_fp32 = flow.cast(y_fp16, dtype=flow.float32)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(y_fp32 * multipler)
return y_fp32
def test_job(x: oft.Numpy.Placeholder(input_shape, dtype=flow.float32)):
v = flow.get_variable(
name="v",
shape=(1, ),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
)
x = x + v
x1 = flow.identity(x)
x2 = flow.identity(x)
flow.watch_diff(x1, test_global_storage.Setter("x1_diff"))
flow.watch_diff(x2, test_global_storage.Setter("x2_diff"))
x1 = flow.cast(x1, data_type)
x2 = flow.cast(x2, data_type)
y1 = flow.layers.batch_normalization_relu(x1, axis=axis, name="BN1")
y2 = flow.math.relu(
flow.layers.batch_normalization(x2, axis=axis, name="BN2"))
y1 = flow.cast(y1, flow.float32)
y2 = flow.cast(y2, flow.float32)
flow.watch(y1, test_global_storage.Setter("y1"))
flow.watch(y2, test_global_storage.Setter("y2"))
y1 = flow.where(flow.math.greater(y2, v), y1, v)
y2 = flow.where(flow.math.greater(y1, v), y2, v)
loss = y1 + y2
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[0.001]),
momentum=0).minimize(flow.math.reduce_sum(loss))
return loss
def MatmulJob():
with flow.scope.placement(device_type, "0:0"):
a = flow.get_variable(
"a",
shape=a_shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=0,
maxval=1),
trainable=True,
)
b = flow.get_variable(
"b",
shape=b_shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=0,
maxval=1),
trainable=True,
)
if data_type == "float16":
out = flow.matmul(
flow.cast(a, dtype=flow.float16),
flow.cast(b, dtype=flow.float16),
transpose_a,
transpose_b,
alpha,
)
c = flow.get_variable(
"c",
shape=out.shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=-1,
maxval=1),
trainable=True,
)
loss = flow.cast(out + flow.cast(c, dtype=flow.float16),
dtype=flow.float)
else:
out = flow.matmul(a, b, transpose_a, transpose_b, alpha)
c = flow.get_variable(
"c",
shape=out.shape,
dtype=dtype,
initializer=flow.random_uniform_initializer(minval=-1,
maxval=1),
trainable=True,
)
loss = out + c
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(loss)
flow.watch(a, test_global_storage.Setter("a"))
flow.watch_diff(a, test_global_storage.Setter("a_diff"))
flow.watch(b, test_global_storage.Setter("b"))
flow.watch_diff(b, test_global_storage.Setter("b_diff"))
flow.watch(c, test_global_storage.Setter("c"))
flow.watch_diff(c, test_global_storage.Setter("c_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
return loss
def RepeatAccJob(a: oft.Numpy.Placeholder(shape)):
if dtype == "float16":
return flow.cast(
flow.acc(flow.repeat(flow.cast(a, flow.float16), acc_num),
acc_num),
flow.float,
)
else:
return flow.acc(flow.repeat(a, acc_num), acc_num)
def ReduceSumLikeJob(x: oft.Numpy.Placeholder(input_shape)):
with flow.scope.placement(device_type, "0:0"):
if data_type == "float16":
x = flow.cast(x, dtype=flow.float16)
like = flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
y = reduce_sum_like(x, like, axis=axis)
y = flow.cast(y, dtype=flow.float32)
else:
like = flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
y = reduce_sum_like(x, like, axis=axis)
return y
def _CreateAttentionMaskFromInputMask(to_mask_blob, from_seq_length, to_seq_length):
output = flow.cast(to_mask_blob, dtype=flow.float)
output = flow.reshape(output, [-1, 1, to_seq_length])
zeros = flow.constant(0.0, dtype=flow.float, shape=[from_seq_length, to_seq_length])
attention_mask_blob = zeros + output
attention_mask_blob = flow.reshape(
attention_mask_blob, [-1, 1, from_seq_length, to_seq_length]
)
attention_mask_blob = flow.cast(attention_mask_blob, dtype=flow.float)
addr_blob = (attention_mask_blob - 1.0) * 10000.0
return addr_blob
def ReduceSumJob(x: oft.Numpy.Placeholder(input_shape)):
with flow.scope.placement(device_type, "0:0"):
if data_type == "float16":
y = flow.cast(
flow.math.reduce_sum(flow.cast(x, dtype=flow.float16),
axis=axis,
keepdims=keepdims),
dtype=flow.float32,
)
else:
y = flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
return y
def _CreateAddrFromAttentionMask(attention_mask_blob, from_seq_length,
to_seq_length):
attention_mask_blob = flow.reshape(attention_mask_blob,
[-1, 1, from_seq_length, to_seq_length])
attention_mask_blob = flow.cast(attention_mask_blob, dtype=flow.float)
addr_blob = (attention_mask_blob - 1.0) * 10000.0
return addr_blob
def oneflow_range_gpu() -> List[tp.Numpy]:
with flow.scope.placement(device_type, machine_ids):
out_1 = flow.range(1,
10,
3,
dtype=flow.float64,
name="range_float64")
out_2 = flow.range(3,
6,
1,
dtype=flow.float32,
name="range_float32")
out_3 = flow.range(4,
13,
4,
dtype=flow.float32,
name="range_float16")
out_3 = flow.cast(out_3, dtype=flow.float32)
out_4 = flow.range(3, dtype=flow.int32, name="range_int32")
out_5 = flow.range(0, 6, 2, dtype=flow.int64, name="range_int64")
x_var = flow.get_variable(
"gpu_input",
shape=(3, ),
dtype=flow.float32,
initializer=flow.constant_initializer(0.0),
)
x_gpu_out = x_var + out_2
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(x_gpu_out)
return [out_1, out_2, out_3, out_4, out_5]
def _CreateAttentionMaskFromInputMask(to_mask_blob, from_seq_length,
to_seq_length):
output = flow.cast(to_mask_blob, dtype=flow.float)
output = flow.reshape(output, [-1, 1, to_seq_length])
zeros = flow.constant(0.0,
dtype=flow.float,
shape=[from_seq_length, to_seq_length])
output = zeros + output
return output
def unsorted_segment_sum_like_job(
data: oft.Numpy.Placeholder(data.shape, dtype=dtype),
segment_ids: oft.Numpy.Placeholder(segment_ids.shape, dtype=flow.int32),
like: oft.Numpy.Placeholder(out_shape, dtype=dtype),
):
with flow.scope.placement(device, "0:0"):
if data_type == "float16":
data = flow.cast(data, dtype=flow.float16)
like = flow.cast(like, dtype=flow.float16)
return flow.cast(
flow.math.unsorted_segment_sum_like(
data=data, segment_ids=segment_ids, like=like, axis=axis
),
dtype=dtype,
)
else:
return flow.math.unsorted_segment_sum_like(
data=data, segment_ids=segment_ids, like=like, axis=axis
)
def oneflow_relu6(of_input_1: tp.Numpy.Placeholder(
shape=input_1.shape, dtype=flow.float32)) -> tp.Numpy:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=input_1.shape,
dtype=flow.float32,
initializer=flow.zeros_initializer(),
name="x_var",
)
x_var = of_input_1 + v
x_f16 = flow.cast(x_var, flow.float16)
of_relu6_out_f16 = flow.nn.relu6(x_f16)
of_relu6_out_f32 = flow.cast(of_relu6_out_f16, flow.float32)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(of_relu6_out_f32)
flow.watch_diff(x_var, assert_prediction_grad)
return of_relu6_out_f32
def op_function(x: tp.Numpy.Placeholder(input.shape, dtype=flow.float32)):
with flow.scope.placement(device_type, "0:0"):
x_var = flow.get_variable(
name="input",
shape=input.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x_var = flow.cast_to_current_logical_view(x_var)
input_x = x_var + x
x_fp32 = flow.cast(input_x, flow.float32)
x_fp16 = flow.cast(input_x, dtype=flow.float16)
y_fp16 = flow.reflection_pad2d(x_fp16, padding)
y_fp32 = flow.cast(y_fp16, dtype=flow.float32)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0]), momentum=0
).minimize(y_fp32)
flow.watch_diff(x_fp32, _compare_diff)
return y_fp32
def FuseCastScaleJob(x: oft.Numpy.Placeholder(
shape, dtype=in_type)) -> Tuple[oft.Numpy, oft.Numpy]:
with flow.scope.placement(device, "0:0-0"):
scale = flow.get_variable(
"scale",
shape=(1, ),
dtype=out_type,
initializer=flow.random_uniform_initializer(),
trainable=False,
)
loss = flow.cast(x, dtype=out_type) * scale
return (loss, scale)
def ActivationJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=shape,
dtype=data_type,
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
x_cast = flow.cast(x, flow.bfloat16)
loss = of_activation_map[activation_type](x_cast)
loss = flow.cast(loss, flow.float)
lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [1e-4])
flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
return loss
def test_job(
x: oft.Numpy.Placeholder(input_shape, dtype=flow.float32),
labels: oft.Numpy.Placeholder(label_shape, dtype=flow.int32),
):
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
name="v",
shape=(1,),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
)
x = x + v
x1 = flow.identity(x)
x2 = flow.identity(x)
flow.watch_diff(x1, test_global_storage.Setter("x1_diff"))
flow.watch_diff(x2, test_global_storage.Setter("x2_diff"))
x1 = flow.cast(x1, data_type)
x2 = flow.cast(x2, data_type)
with flow.scope.placement(device_type, "0:0-3"):
y1 = (
flow.combined_margin_loss(
x1.with_distribute(flow.distribute.split(1)),
labels.with_distribute(flow.distribute.broadcast()),
m1,
m2,
m3,
)
* s
)
y2 = margin_loss(m1, m2, m3, s, x2, labels)
with flow.scope.placement(device_type, "0:0"):
y1 = flow.cast(y1, flow.float)
y2 = flow.cast(y2, flow.float)
flow.watch(y1, test_global_storage.Setter("y1"))
flow.watch(y2, test_global_storage.Setter("y2"))
loss = y1 + y2
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.001]), momentum=0
).minimize(flow.math.reduce_sum(loss))
return loss
def gather_fn(
params_def: oft.Numpy.Placeholder(input.shape, dtype=flow.float32),
indices_def: oft.Numpy.Placeholder(index.shape, dtype=index_type),
) -> oft.Numpy:
with flow.scope.placement(device_type, "0:0"):
x_var = flow.get_variable(
"input",
shape=input.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x_var = flow.cast_to_current_logical_view(x_var)
x = x_var + params_def
x_int32 = flow.cast(x, dtype=flow.int32)
y_int32 = flow.dim_gather(x, dim, indices_def)
y_fp32 = flow.cast(y_int32, dtype=flow.float32)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.001]), momentum=0
).minimize(y_fp32)
flow.watch_diff(x, _compare_diff)
return y_fp32
def FlowJob(x_full_precision: oft.Numpy.Placeholder(x.shape, dtype=dtype)):
with flow.scope.placement(device_type, "0:0"):
x_full_precision += flow.get_variable(
name="v1",
shape=(1, ),
dtype=dtype,
initializer=flow.zeros_initializer())
if data_type == "float16":
x = flow.cast(x_full_precision, flow.float16)
else:
x = x_full_precision
y = flow.layers.batch_normalization(x,
*flow_args,
trainable=trainable,
training=training)
y = flow.cast(y, flow.float)
if trainable:
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(y)
flow.watch_diff(x_full_precision,
test_global_storage.Setter("x_diff"))
return y
def SoftmaxCrossEntropyWithLogitsJob(labels: oft.Numpy.Placeholder(
shape, dtype)):
with flow.scope.placement(device_type, "0:0"):
if data_type == "float16":
x = flow.get_variable(
"x",
shape=shape,
dtype=dtype,
initializer=flow.constant_initializer(20),
trainable=True,
)
loss = flow.cast(
flow.nn.softmax_cross_entropy_with_logits(
flow.cast(labels, dtype=flow.float16),
flow.cast(x, dtype=flow.float16),
),
dtype=flow.float,
)
else:
x = flow.get_variable(
"x",
shape=shape,
dtype=type_name_to_flow_type[data_type],
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
loss = flow.nn.softmax_cross_entropy_with_logits(labels=labels,
logits=x)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
return loss
def FlowJob(
value: oft.Numpy.Placeholder(x_shape), bias: oft.Numpy.Placeholder(bias_shape)
):
with flow.scope.placement(device_type, "0:0"):
value += flow.get_variable(
name="v1",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
bias += flow.get_variable(
name="v2",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x1 = flow.identity(value)
x2 = flow.identity(value)
bias1 = flow.identity(bias)
bias2 = flow.identity(bias)
flow.watch_diff(x1, test_global_storage.Setter("x1_diff"))
flow.watch_diff(x2, test_global_storage.Setter("x2_diff"))
flow.watch_diff(bias1, test_global_storage.Setter("bias1_diff"))
flow.watch_diff(bias2, test_global_storage.Setter("bias2_diff"))
if data_type == "float16":
y1 = flow.cast(
flow.math.gelu(
flow.nn.bias_add(
flow.cast(x1, dtype=flow.float16),
flow.cast(bias1, dtype=flow.float16),
data_format=data_format,
)
),
dtype=flow.float,
)
y2 = flow.cast(
flow.nn.fused_bias_add_gelu(
flow.cast(x2, dtype=flow.float16),
flow.cast(bias2, dtype=flow.float16),
data_format=data_format,
),
dtype=flow.float,
)
else:
y1 = flow.math.gelu(
flow.nn.bias_add(x1, bias1, data_format=data_format)
)
y2 = flow.nn.fused_bias_add_gelu(x2, bias2, data_format=data_format)
flow.watch(y1, test_global_storage.Setter("y1"))
flow.watch(y2, test_global_storage.Setter("y2"))
flow.watch_diff(y1, test_global_storage.Setter("y1_diff"))
flow.watch_diff(y2, test_global_storage.Setter("y2_diff"))
loss = y1 + y2
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.001]), momentum=0
).minimize(flow.math.reduce_sum(loss))
return loss
def op_function(x: tp.Numpy.Placeholder(input.shape, dtype=flow.float32)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="input",
shape=input.shape,
dtype=flow.float32,
initializer=flow.zeros_initializer(),
)
y_int32 = flow.reflection_pad2d(x, padding)
y_fp32 = flow.cast(y_int32, dtype=flow.float32)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0]), momentum=0
).minimize(y_fp32)
flow.watch_diff(x, _compare_diff)
return y_fp32
def CastJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"in",
shape=input_shape,
dtype=type_name_to_flow_type[dtype],
initializer=flow.random_uniform_initializer(),
trainable=True,
)
loss = flow.cast(x, dtype=flow.float)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
return loss
def cast_to_static_shape_fn(
x: flow.typing.ListNumpy.Placeholder(shape=shape, dtype=dtype)
) -> flow.typing.ListNumpy:
x_var = flow.get_variable(
name="x_var",
shape=(1,),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
)
x = x + flow.cast(x_var, dtype=dtype)
y = flow.cast_to_static_shape(x)
test_case.assertFalse(y.is_dynamic)
if require_grad:
flow.watch_diff(x, compare_diff_fn)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.001]), momentum=0
).minimize(y)
return y
