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