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, }