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([], [0.001]), momentum=0
).minimize(z)
flow.watch_diff(x_var, dz_dx_watcher)
flow.watch_diff(y_var, dz_dy_watcher)
return z
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([], [0.0001]), 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 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 PartialFcJob(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-weight",
shape=(num_classes, 128),
dtype=flow.float,
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()
weight_distribute = flow.distribute.split(0)
(
maped_label,
sampled_label,
sampled_weight,
) = flow.distributed_partial_fc_sample(
weight=x.with_distribute(weight_distribute),
label=labels.with_distribute(lebels_distribute),
num_sample=num_sample,
)
with flow.scope.placement(device_type, "0:0"):
sampled_weight = flow.identity(sampled_weight)
loss = flow.math.square(sampled_weight)
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_diff(sampled_weight,
test_global_storage.Setter("sampled_weight_diff"))
return (x, maped_label, sampled_label, sampled_weight)
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([], [0.001]), momentum=0
).minimize(mseloss_mean)
return {
"of_mse_loss": mseloss,
"of_mse_loss_mean": mseloss_mean,
"of_mse_loss_sum": mseloss_sum,
}
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"):
labels = flow.parallel_cast(labels,
distribute=flow.distribute.broadcast())
logits = flow.parallel_cast(
x, distribute=flow.distribute.split(len(x.shape) - 1))
loss = flow.nn.distributed_sparse_softmax_cross_entropy_with_logits(
labels, logits)
loss = flow.math.square(loss)
with flow.scope.placement(device_type, "0:0"):
loss = flow.identity(loss)
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 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,
quantization_formula,
per_layer_quantization,
)
out = flow.quantization.fake_quantization(
input_x,
scale,
zero_point,
quantization_bit,
quantization_scheme,
quantization_formula,
)
loss = flow.math.reduce_mean(out)
flow.optimizer.Adam(
flow.optimizer.PiecewiseConstantScheduler(
[], [0.001])).minimize(loss)
return out
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 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([], [0.001]), momentum=0
).minimize(y)
flow.watch_diff(x_var, dx_watcher)
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(
[], [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 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([], [0.001]), 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 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(
[], [0.001]),
momentum=0).minimize(z)
flow.watch_diff(x, compare_dz_dx)
flow.watch_diff(y, compare_dz_dy)
return z
def ctc_loss_job(
log_probs: tp.Numpy.Placeholder(shape=(max_input_length, batch_size,
num_classes),
dtype=flow_data_type),
targets: tp.Numpy.Placeholder(shape=(batch_size, max_target_length),
dtype=flow.int32),
input_lengths: tp.Numpy.Placeholder(shape=(batch_size, ),
dtype=flow.int32),
target_lengths: tp.Numpy.Placeholder(shape=(batch_size, ),
dtype=flow.int32),
) -> tp.Numpy:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=log_probs.shape,
dtype=flow_data_type,
initializer=flow.zeros_initializer(),
name="x_var",
)
x_var = log_probs + v
flow.watch_diff(x_var, assert_loss_grad)
loss = flow.ctc_loss(
x_var,
targets,
input_lengths,
target_lengths,
blank,
reduction,
zero_infinity,
)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(loss)
return loss
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 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 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([], [0.001]), momentum=0
).minimize(y)
return y
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(
[], [0.001]),
momentum=0).minimize(out)
flow.watch_diff(params_var, params_grad_watcher)
flow.watch_diff(updates_var, updates_grad_watcher)
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,
)
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([], [0.0001]), momentum=0
).minimize(y)
return y
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 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)
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(
[], [0.001]),
momentum=0).minimize(of_Xmum_out)
return of_Xmum_out
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(
[], [0.0001]),
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 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 oneflow_marginloss(
of_anchor: tp.Numpy.Placeholder(shape=anchor.shape),
of_pos: tp.Numpy.Placeholder(shape=pos.shape),
of_neg: tp.Numpy.Placeholder(shape=neg.shape),
) -> Dict[str, tp.Numpy]:
with flow.scope.placement(device_type, "0:0"):
v = flow.get_variable(
shape=anchor.shape,
dtype=flow.float32,
initializer=flow.constant_initializer(0),
name="x_var",
)
x_anchor = of_anchor + v
flow.watch_diff(x_anchor, assert_prediction_grad)
triplet_marginloss = flow.nn.TripletMarginLoss(
x_anchor,
of_pos,
of_neg,
margin=margin,
p=p,
swap=swap,
reduction="none",
name="of_tripletmarginloss",
)
triplet_marginloss_mean = flow.nn.TripletMarginLoss(
x_anchor,
of_pos,
of_neg,
margin=margin,
p=p,
swap=swap,
reduction="mean",
name="of_tripletmarginloss_mean",
)
triplet_marginloss_sum = flow.nn.TripletMarginLoss(
x_anchor,
of_pos,
of_neg,
margin=margin,
p=p,
swap=swap,
reduction="sum",
name="of_tripletmarginloss_sum",
)
with flow.scope.placement(device_type, "0:0"):
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.001]),
momentum=0).minimize(triplet_marginloss_mean)
return {
"of_triplet_margin_loss": triplet_marginloss,
"of_triplet_margin_loss_mean": triplet_marginloss_mean,
"of_triplet_margin_loss_sum": triplet_marginloss_sum,
}
def TrainJob():
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,
)
flow.watch_diff(x, CheckOnes)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler(
[], [0.0001]),
momentum=0).minimize(x)
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 = x + x_blob
y = flow.gather(x, i_blob, axis=axis, batch_dims=batch_dims)
lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.001])
flow.optimizer.SGD(lr_scheduler, momentum=0).minimize(y)
flow.watch_diff(x, compare_fn)
return y
def FlowJob(x: oft.Numpy.Placeholder(x.shape)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="v1",
shape=(1, ),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
loss = flow_op(x, *flow_args)
flow.optimizer.SGD(flow.optimizer.PiecewiseConstantScheduler([],
[0]),
momentum=0).minimize(loss)
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
return loss
def op_function(x: tp.Numpy.Placeholder(input.shape, dtype=value_type)):
with flow.scope.placement(device_type, "0:0"):
x += flow.get_variable(
name="input",
shape=input.shape,
dtype=value_type,
initializer=flow.zeros_initializer(),
)
out = flow.reflection_pad2d(x, padding)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [0]), momentum=0
).minimize(out)
flow.watch_diff(x, _compare_diff)
return out
def watch_matmul_diff_job(
images: tp.Numpy.Placeholder((3, 3), dtype=flow.float)
) -> None:
weight_initializer = flow.constant_initializer(2)
weight_shape = (3, 1)
weight = flow.get_variable(
"three-weight", shape=weight_shape, initializer=weight_initializer
)
weight_broadcast = flow.broadcast_like(
weight, like=images, broadcast_axes=(1,)
)
lr_scheduler = flow.optimizer.PiecewiseConstantScheduler([], [0.1])
flow.optimizer.SGD(lr_scheduler, momentum=0.9).minimize(weight_broadcast)
flow.watch_diff(weight, watch_diff_handler)
