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 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 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 ConcatJob():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"x",
shape=x_shape,
dtype=type_name_to_flow_type[dtype],
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
y = flow.get_variable(
"y",
shape=y_shape,
dtype=type_name_to_flow_type[dtype],
initializer=flow.random_uniform_initializer(minval=-10,
maxval=10),
trainable=True,
)
x = flow.cast_to_current_logical_view(x)
y = flow.cast_to_current_logical_view(y)
loss = flow.concat([x, y], 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(y, test_global_storage.Setter("y"))
flow.watch_diff(y, test_global_storage.Setter("y_diff"))
flow.watch(loss, test_global_storage.Setter("loss"))
flow.watch_diff(loss, test_global_storage.Setter("loss_diff"))
return loss
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.losses.add_loss(z)
flow.watch_diff(x_var, dz_dx_watcher)
flow.watch_diff(y_var, dz_dy_watcher)
return z
def foo_job(input_def: oft.Numpy.Placeholder(shape=(2, 5))):
var = flow.get_variable(
name="var",
shape=(2, 5),
dtype=flow.float,
initializer=flow.ones_initializer(),
)
input_def = flow.cast_to_current_logical_view(input_def)
var = flow.cast_to_current_logical_view(var)
output = var + input_def
return output
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 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.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(y)
return y
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 hybrid_concat_job(
input_0_def: oft.ListNumpy.Placeholder(shape=static_shape, dtype=flow.float),
input_1_def: oft.ListNumpy.Placeholder(shape=static_shape, dtype=flow.float),
):
var = flow.get_variable(
"var",
shape=static_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
trainable=True,
)
constant = flow.constant(1.0, dtype=flow.float, shape=rand_sub_shape)
inputs = [
flow.cast_to_current_logical_view(input)
for input in [var, input_0_def, input_1_def, constant]
]
concated = flow.concat(inputs, axis=axis, max_dim_size=max_dim_size,)
if verbose:
print("concated static shape:", concated.shape)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(concated)
flow.watch_diff(var, compare_var_diff)
if max_dim_size is None:
test_case.assertTrue(
concated.shape[axis] == (static_shape[axis] * 3 + rand_sub_shape[axis])
)
else:
test_case.assertTrue(concated.shape[axis] == max_dim_size)
return var, concated
def dynamic_concat_job(
input_0_def: oft.ListNumpy.Placeholder(
shape=input_static_shape, dtype=flow.float
),
input_1_def: oft.ListNumpy.Placeholder(
shape=input_static_shape, dtype=flow.float
),
):
var_0 = flow.get_variable(
"Var0",
shape=(1,),
dtype=flow.float,
initializer=flow.constant_initializer(value=1, dtype=flow.float),
trainable=True,
)
var_1 = flow.get_variable(
"Var1",
shape=(1,),
dtype=flow.float,
initializer=flow.constant_initializer(value=1, dtype=flow.float),
trainable=True,
)
var_0 = flow.cast_to_current_logical_view(var_0)
var_1 = flow.cast_to_current_logical_view(var_1)
input_0_def = flow.cast_to_current_logical_view(input_0_def)
input_1_def = flow.cast_to_current_logical_view(input_1_def)
if callable(watch_cb):
flow.watch(var_0, watch_cb)
flow.watch(var_1, watch_cb)
flow.watch(flow.identity(input_0_def), watch_cb)
flow.watch(flow.identity(input_1_def), watch_cb)
var_0 = var_0 * input_0_def
var_1 = var_1 * input_1_def
if callable(watch_cb):
flow.watch(var_0, watch_cb)
flow.watch(var_1, watch_cb)
result = flow.concat(
[var_0, var_1], axis=axis, max_dim_size=input_static_shape[axis]
)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), momentum=0
).minimize(result)
flow.watch_diff(var_0, make_watch_diff_cb(0))
flow.watch_diff(var_1, make_watch_diff_cb(1))
return result
def DynamicConvJob(x: oft.ListNumpy.Placeholder((10, 3, 100, 100))):
with flow.scope.placement(device_type, "0:0"):
x_var = flow.get_variable(
name="v1",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
x_var = flow.cast_to_current_logical_view(x_var)
x += x_var
if data_format == "NCHW":
weight_shape = (filters, x_shape[1] // groups, kernel_size,
kernel_size)
else:
weight_shape = (filters, kernel_size, kernel_size,
x_shape[3] // groups)
weight = flow.get_variable(
"conv-weight",
shape=weight_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=0,
maxval=100),
)
weight = flow.cast_to_current_logical_view(weight)
loss = flow.nn.conv2d(
x,
weight,
strides=[stride, stride],
padding=of_padding,
data_format=data_format,
dilations=[1, 1],
groups=groups,
)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-4]),
momentum=0).minimize(loss)
flow.watch(x, global_storage_setter("x"))
flow.watch_diff(x, global_storage_setter("x_diff"))
flow.watch(weight, global_storage_setter("weight"))
flow.watch_diff(weight, global_storage_setter("weight_diff"))
flow.watch(loss, global_storage_setter("loss"))
flow.watch_diff(loss, global_storage_setter("loss_diff"))
return loss
def do_gather_nd(x, index):
x_var = flow.get_variable(
"params", shape=(1,), dtype=x_dtype, initializer=flow.zeros_initializer(),
)
x = x + flow.cast_to_current_logical_view(x_var)
y = flow.gather_nd(x, index)
if need_grad:
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
if callable(comp_diff_fn):
flow.watch_diff(x, comp_diff_fn)
return y
def do_gather(x_blob, i_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + x_blob
y = flow.gather(x, i_blob, axis=axis, batch_dims=batch_dims)
flow.losses.add_loss(y)
flow.watch_diff(x, compare_fn)
return y
def do_scatter_nd(indices_blob, updates_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"updates",
shape=updates.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + updates_blob
y = flow.scatter_nd(indices_blob, x, shape)
flow.losses.add_loss(y)
flow.watch_diff(x, compare_fn)
return y
def DynamicReshapeJob(x: oft.ListNumpy.Placeholder(data_shape)):
reshape_out1 = flow.reshape(x, (-1, 20))
my_model = flow.get_variable(
"my_model",
shape=(20, 32),
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
my_model = flow.cast_to_current_logical_view(my_model)
mm_out = flow.matmul(reshape_out1, my_model)
reshape_out2 = flow.reshape(mm_out, (-1, 8, 4))
flow.losses.add_loss(reshape_out2)
return reshape_out1
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.losses.add_loss(y)
flow.watch_diff(x, grad_cb)
return y
def oneflow_Xmum(
of_input_1: tp.ListNumpy.Placeholder(shape=data_shape),
of_input_2: tp.ListNumpy.Placeholder(shape=data_shape),
) -> tp.ListNumpy:
with flow.scope.placement(device_type, "0:0"):
v1 = flow.get_variable(
shape=(1, ),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
name="x1_var",
)
v1 = flow.cast_to_current_logical_view(v1)
x1_var = of_input_1 + v1
v2 = flow.get_variable(
shape=(1, ),
dtype=flow.float32,
initializer=flow.zeros_initializer(),
name="x2_var",
)
v2 = flow.cast_to_current_logical_view(v2)
x2_var = of_input_2 + v2
flow.watch_diff(x1_var,
assert_prediction_grad) # Only Compare input1 Grad
if compare_type == "maximum":
of_Xmum_out = flow.math.maximum(x1_var, x2_var)
elif compare_type == "minimum":
of_Xmum_out = flow.math.minimum(x1_var, x2_var)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(of_Xmum_out)
return of_Xmum_out
def do_gather(x_blob, i_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"params",
shape=params.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + x_blob
y = flow.gather(x, i_blob, axis=axis, batch_dims=batch_dims)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [1e-3]),
momentum=0).minimize(y)
flow.watch_diff(x, compare_fn)
return y
def do_scatter_nd(indices_blob, updates_blob):
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"updates",
shape=updates.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x = flow.cast_to_current_logical_view(x)
x = x + updates_blob
y = flow.scatter_nd(indices_blob, x, shape)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(x, compare_fn)
return y
def DynamicReshapeJob(x: oft.ListNumpy.Placeholder(data_shape)):
reshape_out1 = flow.reshape(x, (-1, 20))
my_model = flow.get_variable(
"my_model",
shape=(20, 32),
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=-10, maxval=10),
trainable=True,
)
my_model = flow.cast_to_current_logical_view(my_model)
mm_out = flow.matmul(reshape_out1, my_model)
reshape_out2 = flow.reshape(mm_out, (-1, 8, 4))
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), momentum=0
).minimize(reshape_out2)
return reshape_out1
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 reshape(
x: remote_blob_util.BlobDef, shape: Sequence[int], name: Optional[str] = None
) -> remote_blob_util.BlobDef:
r"""Reshapes a blob.
Args:
x: A `Blob`.
shape: Shape of the output blob.
name: A name for the operation (optional).
Returns:
A `Blob`, has the same type as `x`.
"""
x = flow.cast_to_current_logical_view(x)
assert isinstance(shape, tuple) or isinstance(shape, list)
shape = list(shape)
assert all(dim == -1 or dim > 0 for dim in shape)
assert shape.count(-1) <= 1
if not x.is_dynamic:
if name is None:
name = id_util.UniqueStr("Reshape_")
return (
flow.user_op_builder(name)
.Op("reshape")
.Input("in", [x])
.Output("out")
.Attr("shape", infer_shape(x, shape))
.Build()
.InferAndTryRun()
.RemoteBlobList()[0]
)
else:
op_conf = op_conf_util.OperatorConf()
setattr(
op_conf,
"name",
name if name is not None else id_util.UniqueStr("DynamicReshape_"),
)
setattr(op_conf.dynamic_reshape_conf, "in", x.unique_name)
op_conf.dynamic_reshape_conf.shape.dim.extend(list(shape))
setattr(op_conf.dynamic_reshape_conf, "out", "out")
interpret_util.Forward(op_conf)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_conf.name
lbi.blob_name = "out"
return remote_blob_util.RemoteBlob(lbi)
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 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.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
).minimize(y)
flow.watch_diff(params_var, compare_fn)
return y
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 DynamicDimGatherJob(
params_def: oft.ListNumpy.Placeholder(datashape, dtype=flow.float32),
index_def: oft.ListNumpy.Placeholder(datashape, dtype=flow.int32),
) -> oft.ListNumpy:
x_var = flow.get_variable(
"input",
shape=(1, ),
dtype=flow.float32,
initializer=flow.constant_initializer(0),
)
x_var = flow.cast_to_current_logical_view(x_var)
x = x_var + params_def
y = flow.dim_gather(x, sample["dim"], index_def)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[1e-3]),
momentum=0).minimize(y)
flow.watch_diff(x, _compare_diff)
return y
