def clip(values_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"values",
shape=values.shape,
dtype=data_type,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + values_blob
y = flow.clip_by_value(x, min, max)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(x, grad_cb)
return y
def fake_quantization():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
return flow.quantization.fake_quantization(
x,
*flow.quantization.min_max_observer(
x,
per_layer_quantization=per_layer,
quantization_scheme=scheme,
),
quantization_scheme=scheme,
)
def ReduceMaxJob(x: oft.Numpy.Placeholder(input_shape, dtype=flow.float)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="v1",
shape=input_shape,
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
loss = flow.math.reduce_max(x, axis=axis, keepdims=keepdims)
loss = flow.identity(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 PolyValJob(x: tp.Numpy.Placeholder(shape=in_shape)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="x",
shape=in_shape,
dtype=flow_data_type,
initializer=flow.zeros_initializer(),
trainable=True,
)
flow.watch_diff(x, assert_prediction_grad)
out = flow.math.polyval(coeffs, x)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(out)
return out
def testSGD(random_mask: flow.typing.Numpy.Placeholder(
x_shape, dtype=flow.float32)) -> flow.typing.Numpy:
with flow.scope.placement(device_type, "0:0-0"):
x = flow.get_variable(
name="x",
shape=x_shape,
dtype=flow.float32,
initializer=flow.random_uniform_initializer(minval=0,
maxval=100),
trainable=True,
)
loss = flow.math.reduce_mean(x * random_mask)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [learning_rate]),
momentum=momentum,
).minimize(loss)
return x
def ReshapeJob():
with flow.scope.placement(device_type, "0:0-{}".format(device_num - 1)):
x = flow.get_variable(
"var_x",
shape=input_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
distribute=flow.distribute.split(2),
)
loss = flow.reshape(x, shape)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), momentum=0
).minimize(loss)
return x, loss
def do_unsorted_segment_sum(x_blob, i_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"data",
shape=data.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = x + x_blob
y = flow.math.unsorted_segment_sum(
x, i_blob, axis=axis, num_segments=num_segments
)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(x, compare_fn)
return y
def static_concat_job(
input_0_def: oft.Numpy.Placeholder(shape=shape, dtype=flow.float),
input_1_def: oft.Numpy.Placeholder(shape=shape, dtype=flow.float),
):
var = flow.get_variable(
"var",
shape=shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
trainable=True,
)
concated = flow.concat([input_0_def, input_1_def, var], axis=axis)
test_case.assertTrue(not concated.is_dynamic)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[1e-3]),
momentum=0).minimize(concated)
flow.watch_diff(var, compare_var_diff)
return var, concated
def SquareJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
loss = flow.math.square(x)
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 ReduceSumJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"in",
shape=input_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
)
loss = flow.math.reduce_sum(x, axis=axis, keepdims=keepdims)
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 oneflow_ones() -> tp.Numpy:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=np_out_ones.shape,
dtype=flow.float32,
initializer=flow.zeros_initializer(),
name="x_var",
)
of_ones = flow.ones(shape=input_shape, dtype=flow.float32)
of_out = of_ones + v
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(of_out)
return of_ones
def gather_nd_fn(
params_def: oft.ListNumpy.Placeholder(static_params_shape, dtype=flow.float),
indices_def: oft.ListNumpy.Placeholder(indices.shape, dtype=flow.int32),
):
with flow.scope.placement("gpu", "0:0"):
one_var = flow.get_variable(
"one",
shape=(1,),
dtype=flow.float32,
initializer=flow.constant_initializer(1),
)
one_var = flow.cast_to_current_logical_view(one_var)
params_var = params_def * one_var
y = flow.gather_nd(params_var, indices_def)
flow.losses.add_loss(y)
flow.watch_diff(params_var, compare_fn)
return y
def LeakyReluJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=type_name_to_flow_type[data_type],
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
loss = flow.nn.leaky_relu(x, alpha=alpha)
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 L2NormalizeJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=type_name_to_flow_type[data_type],
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
loss = flow.math.l2_normalize(x, axis=axis, epsilon=epsilon)
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 slice_with_grad_job(
x: otp.Numpy.Placeholder(shape=input_shape, dtype=dtype)
) -> otp.Numpy:
var = flow.get_variable(
shape=input_shape,
dtype=dtype,
initializer=flow.constant_initializer(0.0),
name="variable",
)
x = x + var
if callable(watch_diff_cb):
flow.watch_diff(x, watch_diff_cb)
y = flow.slice_v2(x, slice_tup_list, name="SliceWithGrad")
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
return y
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("gpu", "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("gpu", "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("gpu", "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 testIndexedSlicesSGD(
sparse_ids: flow.typing.Numpy.Placeholder(ids_shape, dtype=flow.int32),
) -> flow.typing.Numpy:
with flow.scope.placement(device_type, "0:0"):
embedding_table = flow.get_variable(
name="embeddings",
shape=model_shape,
initializer=flow.random_uniform_initializer(minval=0,
maxval=100),
)
embedding = flow.gather(params=embedding_table * mul_scalar,
indices=sparse_ids)
loss = flow.math.reduce_mean(embedding)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [learning_rate]),
momentum=momentum_beta,
).minimize(loss)
return embedding_table
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,
)
loss = of_activation_map[activation_type](x)
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 FlattenJob() -> flow.typing.Numpy:
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"in",
shape=input_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
)
loss = flow.flatten(x, start_dim=start_dim, end_dim=end_dim)
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"))
return loss
def TestJob(
prediction: oft.Numpy.Placeholder(
prediction_shape, dtype=type_name_to_flow_type[data_type]),
label: oft.Numpy.Placeholder(
prediction_shape, dtype=type_name_to_flow_type[data_type]),
):
v = flow.get_variable(
"prediction",
shape=prediction_shape,
dtype=type_name_to_flow_type[data_type],
initializer=flow.constant_initializer(0),
trainable=True,
)
flow.watch_diff(v, assert_prediction_grad)
prediction += v
with flow.scope.placement(device_type, "0:0"):
loss = flow.smooth_l1_loss(prediction, label, beta)
flow.losses.add_loss(loss)
return loss
def oneflow_range() -> 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(3, dtype=flow.int32, name="range_int32")
out_4 = flow.range(0, 6, 2, dtype=flow.int64, name="range_int64")
x_var = flow.get_variable(
"cpu_input",
shape=(3,),
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x_out = out_2 + x_var
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(x_out)
return [out_1, out_2, out_3, out_4]
def TransposeJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"input",
shape=input_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
)
loss = flow.transpose(x, perm)
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 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(
[], [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.losses.add_loss(y)
flow.watch_diff(x_full_precision, test_global_storage.Setter("x_diff"))
return y
def ReduceMinJob(x: oft.Numpy.Placeholder(input_shape, dtype=flow.float)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="v1",
shape=input_shape,
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
loss = flow.math.reduce_min(x, axis=axis, keepdims=keepdims)
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 unsorted_batch_segment_sum_job(
data: oft.Numpy.Placeholder(data.shape, dtype=flow.float),
segment_ids: oft.Numpy.Placeholder(segment_ids.shape, dtype=flow.int32),
):
with flow.scope.placement(device, "0:0"):
x = flow.get_variable(
"data",
shape=data.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
data = x + data
res = flow.math.unsorted_batch_segment_sum(
data=data, segment_ids=segment_ids, num_segments=num_segments
)
flow.losses.add_loss(res)
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
flow.watch_diff(res, test_global_storage.Setter("loss_diff"))
return res
def Build_EmbeddingLayer(self,
vocab_size,
embedding_size=128,
word_embedding_name="Embedding_Layer"):
"""
Build a Embedding Layer
:param input_ids_blob:The input ID Blob
:param vocab_size: The input Vocab size
:param embedding_size: The embedding Size
:param initializer_range: The range of Initializer, Use flow.truncated_normal
:param word_embedding_name: The name of Embedding variable
:return: The output and the Embedding table.
"""
self.embedding_table = flow.get_variable(
name=word_embedding_name + "_Embed",
shape=[vocab_size, embedding_size],
dtype=flow.float32,
initializer=flow.random_normal_initializer(0,
self.hidden_size**-0.5))
def two_stage_reduce_job(x: oft.Numpy.Placeholder((4, 20, 20, 20))):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="v1",
shape=(1,),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
with flow.scope.placement(device_type, "0:0-3"):
loss = flow_func(
x.with_distribute(flow.distribute.split(split_axis)),
axis=axis,
keepdims=True,
)
loss = flow.identity(loss)
flow.losses.add_loss(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
return loss
def _default_initializer_for_determining(undetermined_tensor):
assert not undetermined_tensor.is_consistent
variable_name = id_util.UniqueStr("tensor_")
blob = flow.get_variable(
name=variable_name,
shape=tuple(undetermined_tensor.shape),
dtype=undetermined_tensor.dtype,
initializer=undetermined_tensor.data_initializer,
)
determined_tensor = oneflow_api.LocalTensor(
undetermined_tensor.shape,
undetermined_tensor.dtype,
undetermined_tensor.device,
undetermined_tensor.is_lazy,
undetermined_tensor.requires_grad,
True,
undetermined_tensor.retain_grad,
)
determined_tensor._set_blob_object(blob.blob_object)
return determined_tensor, variable_name
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
