def _test_tril_fw_bw(test_case, device, shape, type_name, diagonal, fill_value):
flow.clear_default_session()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
if type_name == "float16":
flow_type = flow.float
np_type = np.float32
else:
flow_type = type_name_to_flow_type[type_name]
np_type = type_name_to_np_type[type_name]
@flow.global_function(type="train", function_config=func_config)
def test_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.tril(flow.cast(x, flow.float16), diagonal), flow.float
)
else:
out = flow.math.tril(x, diagonal)
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
x = np.random.randint(low=0, high=100, size=shape)
test_tril_fw_bw_job(x.astype(np_type)).get()
np_out = np.where(
np.tril(np.ones(shape), diagonal),
test_global_storage.Get("x"),
np.full(shape, fill_value).astype(np_type),
)
np_x_diff = np.tril(test_global_storage.Get("out_diff"), diagonal)
if type_name == "float16":
tolerance = 0.001
else:
tolerance = 1e-05
test_case.assertTrue(
np.allclose(
np_out, test_global_storage.Get("out"), rtol=tolerance, atol=tolerance
)
)
test_case.assertTrue(
np.allclose(
np_x_diff, test_global_storage.Get("x_diff"), rtol=tolerance, atol=tolerance
)
)
def _test_element_wise_mul_fw_bw(test_case, device, shape, type_name):
flow.clear_default_session()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
np_type = type_name_to_np_type[type_name]
flow_type = type_name_to_flow_type[type_name]
@flow.global_function(type="train", function_config=func_config)
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
x = np.random.randint(low=0, high=10, size=shape).astype(np.float32)
y = np.random.randint(low=0, high=10, size=shape).astype(np.float32)
test_element_wise_mul_job(x, y).get()
test_case.assertTrue(
np.allclose(
test_global_storage.Get("x") * test_global_storage.Get("y"),
test_global_storage.Get("out"),
))
test_case.assertTrue(
np.allclose(
test_global_storage.Get("out_diff") * test_global_storage.Get("x"),
test_global_storage.Get("y_diff"),
))
test_case.assertTrue(
np.allclose(
test_global_storage.Get("out_diff") * test_global_storage.Get("y"),
test_global_storage.Get("x_diff"),
))
def _test_masked_fill_fw_bw(test_case,
device,
x_shape,
mask_shape,
type_name,
value=0):
flow.clear_default_session()
func_config = flow.FunctionConfig()
if type_name == "float16":
flow_type = flow.float
np_type = np.float32
else:
flow_type = type_name_to_flow_type[type_name]
np_type = type_name_to_np_type[type_name]
func_config.default_data_type(flow_type)
@flow.global_function(type="train", function_config=func_config)
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
x = np.random.randint(low=0, high=100, size=x_shape)
mask = np.random.randint(low=0, high=2, size=mask_shape)
test_masked_fill_fw_bw_job(x.astype(np_type), mask.astype(np_type)).get()
out_diff = test_global_storage.Get("out_diff")
np_out, np_x_diff = _masked_fill_np_fw_bw(x, mask, out_diff, np_type,
value)
if type_name == "float16":
tolerance = 1e-3
else:
tolerance = 1e-5
test_case.assertTrue(
np.allclose(np_out,
test_global_storage.Get("out"),
rtol=tolerance,
atol=tolerance))
test_case.assertTrue(
np.allclose(np_x_diff,
test_global_storage.Get("x_diff"),
rtol=tolerance,
atol=tolerance))
def _test_element_wise_mul_fw_bw(test_case, device, shape, type_name):
flow.clear_default_session()
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
func_config.train.primary_lr(1e-4)
func_config.train.model_update_conf(dict(naive_conf={}))
np_type = type_name_to_np_type[type_name]
flow_type = type_name_to_flow_type[type_name]
@flow.global_function(func_config)
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.losses.add_loss(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
check_point = flow.train.CheckPoint()
check_point.init()
x = np.random.randint(low=0, high=10, size=shape).astype(np.float32)
y = np.random.randint(low=0, high=10, size=shape).astype(np.float32)
test_element_wise_mul_job(x, y).get()
test_case.assertTrue(
np.allclose(
test_global_storage.Get("x") * test_global_storage.Get("y"),
test_global_storage.Get("out"),
)
)
test_case.assertTrue(
np.allclose(
test_global_storage.Get("out_diff") * test_global_storage.Get("x"),
test_global_storage.Get("y_diff"),
)
)
test_case.assertTrue(
np.allclose(
test_global_storage.Get("out_diff") * test_global_storage.Get("y"),
test_global_storage.Get("x_diff"),
)
)