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_f16 = flow.cast(x, flow.float16)
y_f16 = flow.dim_gather(x_f16, dim, indices_def)
x_f32 = flow.cast(x, flow.float32)
y_f32 = flow.cast(y_f16, flow.float32)
y = flow.dim_gather(x, dim, indices_def)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(y_f32)
flow.watch_diff(x_f32, _compare_diff)
return y_f32
def SparseSoftmaxCrossEntropyWithLogitsJob(labels: oft.Numpy.Placeholder(
(batch_size, ), dtype=type_name_to_flow_type[label_type])):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=(batch_size, num_classes),
dtype=type_name_to_flow_type[data_type],
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
with flow.scope.placement(device_type, "0:0-3"):
lebels_distribute = flow.distribute.broadcast()
logits_distribute = flow.distribute.split(len(x.shape) - 1)
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels.with_distribute(lebels_distribute),
logits=x.with_distribute(logits_distribute),
)
loss = flow.math.square(loss)
with flow.scope.placement(device_type, "0:0"):
loss = flow.identity(loss)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
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 scatter_nd_update_grad_fn(
x_def: oft.Numpy.Placeholder(params.shape, dtype=flow.float),
indices_def: oft.Numpy.Placeholder(indices.shape,
dtype=flow.int32),
y_def: oft.Numpy.Placeholder(updates.shape, dtype=flow.float),
):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
y = flow.get_variable(
"updates",
shape=updates.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = x + x_def
y = y + y_def
z = flow.tensor_scatter_nd_update(x, indices_def, y)
flow.losses.add_loss(z)
flow.watch_diff(x, compare_dz_dx)
flow.watch_diff(y, compare_dz_dy)
return z
def do_tensor_scatter_nd_add(params_blob, indices_blob, updates_blob):
with flow.scope.placement(device_type, "0:0"):
params_var = flow.get_variable(
"params",
shape=params_blob.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
updates_var = flow.get_variable(
"updates",
shape=updates_blob.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
params_var = flow.cast_to_current_logical_view(params_var)
params_blob = flow.cast_to_current_logical_view(params_blob)
updates_blob = flow.cast_to_current_logical_view(updates_blob)
updates_var = flow.cast_to_current_logical_view(updates_var)
params_var = params_var + params_blob
updates_var = updates_var + updates_blob
out = flow.tensor_scatter_nd_add(params_var, indices_blob, updates_var)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(out)
flow.watch_diff(params_var, params_grad_watcher)
flow.watch_diff(updates_var, updates_grad_watcher)
return out
def TestMultiInputJob():
with flow.scope.placement("gpu", "0:0"):
x1 = flow.get_variable(
"x1",
shape=shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
x2 = flow.get_variable(
"x2",
shape=shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
loss = TestMultiInput(x1, x2)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(loss)
flow.watch(x1, test_global_storage.Setter("x1"))
flow.watch_diff(x1, test_global_storage.Setter("x1_diff"))
flow.watch(x2, test_global_storage.Setter("x2"))
flow.watch_diff(x2, test_global_storage.Setter("x2_diff"))
return loss
def instanceNormJob(
of_input: tp.Numpy.Placeholder(shape=input.shape),
multipler: tp.Numpy.Placeholder(shape=input.shape),
) -> 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",
)
x_var = of_input + v
# watch the gradient
flow.watch_diff(x_var, assert_prediction_grad)
if len(of_input.shape) == 3:
out = flow.nn.InstanceNorm1d(x_var, eps=eps, affine=affine)
elif len(of_input.shape) == 4:
out = flow.nn.InstanceNorm2d(x_var, eps=eps, affine=affine)
else: # len(of_input.shape) == 5
out = flow.nn.InstanceNorm3d(x_var, eps=eps, affine=affine)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(out * multipler)
return out
def oneflow_hardtanh(
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_hardtanh_out_f16 = flow.nn.hardtanh(x_f16, min_val, max_val)
of_hardtanh_out_f32 = flow.cast(of_hardtanh_out_f16, flow.float32)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(of_hardtanh_out_f32)
flow.watch_diff(x_var, assert_prediction_grad)
return of_hardtanh_out_f32
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 pooling_job(x: tensor_def(x_shape, dtype=dtype)):
v = flow.get_variable(
"x",
shape=x_shape,
dtype=dtype,
initializer=flow.constant_initializer(0),
trainable=True,
)
v = flow.cast_to_current_logical_view(v)
flow.watch_diff(v, assert_grad)
x += v
with flow.scope.placement(device_type, "0:0"):
pooling_f = None
if pooling_type == "AVG":
pooling_f = getattr(flow.nn, "avg_pool{}d".format(dim))
elif pooling_type == "MAX":
pooling_f = getattr(flow.nn, "max_pool{}d".format(dim))
else:
raise ValueError("pooling_type must be AVG or MAX")
y = pooling_f(
x,
ksize=ksize,
strides=strides,
padding=padding,
data_format=data_format,
)
flow.losses.add_loss(y)
return y
def broadcast_to_compatible_with_fn(x_def: oft.Numpy.Placeholder(
x.shape, dtype=flow.float)):
x_var = flow.get_variable(
"x_var",
shape=x.shape,
dtype=flow.float,
initializer=flow.constant_initializer(0),
trainable=True,
)
compatible_var = [
flow.get_variable(
"compatible_var_{}".format(i),
shape=cp_shape,
dtype=flow.float,
initializer=flow.random_normal_initializer(),
trainable=False,
) for i, cp_shape in enumerate(compatible_shape)
]
x_var = x_var + x_def
y = flow.broadcast_to_compatible_with(x_var, compatible_var)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[1e-3]),
momentum=0).minimize(y)
flow.watch_diff(x_var, dx_watcher)
return y
def do_where(condition, x, y):
with flow.scope.placement(device_type, "0:0"):
x_var = flow.get_variable(
"x",
shape=x.shape,
dtype=flow.float,
initializer=flow.constant_initializer(0),
)
x_var = flow.cast_to_current_logical_view(x_var)
x_var = x_var + x
y_var = flow.get_variable(
"y",
shape=y.shape,
dtype=flow.float,
initializer=flow.constant_initializer(0),
)
y_var = flow.cast_to_current_logical_view(y_var)
y_var = y_var + y
z = flow.where(condition, x_var, y_var)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(z)
flow.watch_diff(x_var, dz_dx_watcher)
flow.watch_diff(y_var, dz_dy_watcher)
return z
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"))
loss = flow.math.reduce_mean(y1 + y2)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0.001]), momentum=0
).minimize(flow.math.reduce_sum(loss))
return loss
def SparseSoftmaxCrossEntropyWithLogitsJob(labels: oft.Numpy.Placeholder(
(batch_size, ), dtype=type_name_to_flow_type[label_type])):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=(batch_size, num_classes),
dtype=type_name_to_flow_type[data_type],
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
prediction = flow.nn.softmax(logits=x)
with flow.scope.placement(device_type, "0:0-3"):
lebels_distribute = flow.distribute.broadcast()
prediction_distribute = flow.distribute.split(
len(prediction.shape) - 1)
loss = flow.nn.sparse_cross_entropy(
labels=labels.with_distribute(lebels_distribute),
prediction=prediction.with_distribute(prediction_distribute),
)
with flow.scope.placement(device_type, "0:0"):
loss = flow.math.square(loss)
flow.losses.add_loss(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 QuantizeJob(input: oft.Numpy.Placeholder(
in_shape, dtype=type_name_to_flow_type[dtype])):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=in_shape,
dtype=input.dtype,
initializer=flow.zeros_initializer(input.dtype),
trainable=True,
)
input_x = input + x
flow.watch_diff(input_x, test_global_storage.Setter("input_diff"))
with flow.scope.placement(device_type, "0:0-%d" % (device_num - 1)):
scale, zero_point = flow.quantization.min_max_observer(
input_x, quantization_bit, quantization_scheme,
per_layer_quantization)
out = flow.quantization.fake_quantization(input_x, scale,
zero_point,
quantization_bit,
quantization_scheme)
loss = flow.math.reduce_mean(out)
flow.optimizer.Adam(
flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]), ).minimize(loss)
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,
seed=seed,
name="dropout"),
dtype,
)
else:
of_out = flow.nn.dropout(x,
rate=rate,
seed=seed,
name="dropout")
loss = flow.math.square(of_out)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
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 slice_update_train_job(
x: otp.Numpy.Placeholder(shape=input_shape, dtype=dtype),
update: otp.Numpy.Placeholder(shape=update_shape, dtype=dtype),
) -> otp.Numpy:
x_var = flow.get_variable(
shape=input_shape,
dtype=dtype,
initializer=flow.constant_initializer(0.0),
name="x",
)
update_var = flow.get_variable(
shape=update_shape,
dtype=dtype,
initializer=flow.constant_initializer(0.0),
name="update",
)
x = x + x_var
update = update + update_var
if callable(diff_watcher_maker):
flow.watch_diff(x, diff_watcher_maker(input_shape))
flow.watch_diff(update, diff_watcher_maker(update_shape))
y = flow.slice_update(x, update, slice_tup_list)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[1e-3]),
momentum=0).minimize(y)
return y
def UpsampleJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"input",
shape=input_shape,
dtype=type_name_to_flow_type[dtype],
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
)
loss = flow.layers.upsample_2d(
x,
size=size,
data_format=data_format,
interpolation=interpolation,
align_corners=align_corners,
)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), 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(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 FlowJob(
x: oft.Numpy.Placeholder(x.shape, dtype=flow_type),
y: oft.Numpy.Placeholder(y.shape, dtype=flow_type),
):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="x",
shape=x.shape,
dtype=flow_type,
initializer=flow.zeros_initializer(),
trainable=True,
)
y += flow.get_variable(
name="y",
shape=y.shape,
dtype=flow_type,
initializer=flow.zeros_initializer(),
trainable=True,
)
loss = flow_op(x, y)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(loss)
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
return loss
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 scatter_nd_update_grad_fn(
x_def: oft.Numpy.Placeholder(params.shape, dtype=flow.float),
indices_def: oft.Numpy.Placeholder(indices.shape, dtype=flow.int32),
y_def: oft.Numpy.Placeholder(updates.shape, dtype=flow.float),
):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
y = flow.get_variable(
"updates",
shape=updates.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = x + x_def
y = y + y_def
z = flow.tensor_scatter_nd_update(x, indices_def, y)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(z)
flow.watch_diff(x, compare_dz_dx)
flow.watch_diff(y, compare_dz_dy)
return z
def SplitLikeJob(x: oft.Numpy.Placeholder(x_shape, dtype=flow.float)):
v = flow.get_variable(
"x",
shape=x_shape,
dtype=flow.float,
initializer=flow.constant_initializer(0),
trainable=True,
)
x += v
like0 = flow.constant(0, dtype=flow.float, shape=like0_shape)
like1 = flow.constant(0, dtype=flow.float, shape=like1_shape)
with flow.scope.placement("gpu", "0:0"):
y0, y1 = split_like(x, [like0, like1], "split_like")
loss = y0
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), 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 y0, y1
def diag_job(
input_tensor: tp.Numpy.Placeholder(shape=(input_shape),
dtype=flow.float),
) -> tp.Numpy:
input_var = flow.get_variable(
"input_tensor",
shape=(input_shape),
dtype=flow.float,
initializer=flow.zeros_initializer(),
trainable=True,
)
input_tensor = input_tensor + input_var
input_tensor = flow.cast_to_current_logical_view(input_tensor)
input_tensor = flow.cast(input_tensor, type_name_to_flow_type[dtype])
output = flow.diag(input_tensor, dim)
if (output.dtype == flow.int64 or output.dtype == flow.int8
or output.dtype == flow.int32):
output = flow.cast(output, flow.float)
flow.optimizer.Adam(
flow.optimizer.PiecewiseConstantScheduler([],
[1e-4])).minimize(output)
flow.watch(input_tensor, test_global_storage.Setter("x"))
flow.watch_diff(input_tensor, test_global_storage.Setter("x_diff"))
flow.watch(output, test_global_storage.Setter("output"))
flow.watch_diff(output, test_global_storage.Setter("output_diff"))
return output
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=-0.1,
maxval=0.1),
trainable=True,
)
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.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
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_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(
[], [1e-4]),
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 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(
[], [1e-4]),
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 pooling_job(x: tensor_def(x_shape, dtype=dtype)):
v = flow.get_variable(
"x",
shape=x_shape,
dtype=dtype,
initializer=flow.constant_initializer(0),
trainable=True,
)
flow.watch_diff(v, assert_grad)
x += v
with flow.scope.placement(device_type, "0:0"):
pooling_f = None
if pooling_type == "AVG":
pooling_f = getattr(flow.nn, "avg_pool{}d".format(dim))
elif pooling_type == "MAX":
pooling_f = getattr(flow.nn, "max_pool{}d".format(dim))
else:
raise ValueError("pooling_type must be AVG or MAX")
padding = pool_conf["padding"]
if padding == "SAME":
padding = "SAME_UPPER"
y = pooling_f(
x,
ksize=ksize,
strides=strides,
padding=padding,
data_format=data_format,
)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(y)
return y
def FlowJob(
x: oft.Numpy.Placeholder(x.shape, dtype=flow_type),
y: oft.Numpy.Placeholder(y.shape, dtype=flow_type),
):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="x",
shape=x.shape,
dtype=flow_type,
initializer=flow.zeros_initializer(),
trainable=True,
)
y += flow.get_variable(
name="y",
shape=y.shape,
dtype=flow_type,
initializer=flow.zeros_initializer(),
trainable=True,
)
loss = flow_op(x, y)
flow.losses.add_loss(loss)
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
return loss
def do_tensor_scatter_nd_add(params_blob, indices_blob, updates_blob):
with flow.scope.placement(device_type, "0:0"):
params_var = flow.get_variable(
"params",
shape=params_blob.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
updates_var = flow.get_variable(
"updates",
shape=updates_blob.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
params_var = flow.cast_to_current_logical_view(params_var)
params_blob = flow.cast_to_current_logical_view(params_blob)
updates_blob = flow.cast_to_current_logical_view(updates_blob)
updates_var = flow.cast_to_current_logical_view(updates_var)
params_var = params_var + params_blob
updates_var = updates_var + updates_blob
out = flow.tensor_scatter_nd_add(params_var, indices_blob,
updates_var)
flow.losses.add_loss(out)
flow.watch_diff(params_var, params_grad_watcher)
flow.watch_diff(updates_var, updates_grad_watcher)
return out
def oneflow_mseloss(
of_input: tp.Numpy.Placeholder(shape=input.shape),
of_target: tp.Numpy.Placeholder(shape=target.shape),
) -> Dict[str, tp.Numpy]:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=input.shape,
dtype=flow.float32,
initializer=flow.zeros_initializer(),
name="x_var",
)
x_var = of_input + v
flow.watch_diff(x_var, assert_prediction_grad)
mseloss = flow.nn.MSELoss(x_var, of_target, reduction="none", name="of_mseloss")
mseloss_mean = flow.nn.MSELoss(
x_var, of_target, reduction="mean", name="of_mseloss_reduce_mean"
)
mseloss_sum = flow.nn.MSELoss(
x_var, of_target, reduction="sum", name="of_mseloss_reduce_sum"
)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(mseloss_mean)
return {
"of_mse_loss": mseloss,
"of_mse_loss_mean": mseloss_mean,
"of_mse_loss_sum": mseloss_sum,
}
