def my_model_fn(mode): self.assertEqual(mode, model_fn_lib.ModeKeys.EVAL) def stage1(features): w1 = variable_scope.get_variable("w1", initializer=1.0) partial = w1 * features return partial def stage2(partial): w2 = variable_scope.get_variable("w2", initializer=1.0) partial = w2 * partial return partial def stage3(partial): squared = partial * partial return partial, squared def eval_metrics_fn(partial, squared): return { "mean": metrics_impl.mean(partial), "loss": squared, } return IPUPipelineEstimatorSpec( mode, computational_stages=[stage1, stage2, stage3], device_mapping=[0, 1, 0], eval_metrics_fn=eval_metrics_fn, gradient_accumulation_count=gradient_accumulation_count)
def my_model_fn(mode): def stage1(global_step, features, labels): w = variable_scope.get_variable(name="w", initializer=initial_w) partial = w * features return global_step, partial, labels def stage2(global_step, partial, labels): loss = partial + labels return global_step, loss def optimizer_function(global_step, loss): lr = 0.1 - 0.001 * global_step opt = gradient_descent.GradientDescentOptimizer(lr) return pipelining_ops.OptimizerFunctionOutput(opt, loss) def eval_metrics_fn(global_step, loss): return { "global_step_observed": metrics_impl.mean(global_step), "loss": loss, } global_step_input = math_ops.cast(training_util.get_global_step(), dtype=np.float32) return IPUPipelineEstimatorSpec( mode, computational_stages=[stage1, stage2], optimizer_function=optimizer_function, eval_metrics_fn=eval_metrics_fn, inputs=[global_step_input], gradient_accumulation_count=gradient_accumulation_count)
def my_model_fn(mode): def stage1(features): w1 = variable_scope.get_variable("w1", initializer=1.0) partial = w1 * features return partial def stage2(partial): squared = partial * partial if arg_type is list: return [partial, squared] assert arg_type is dict # Pass in reverse order just to check that they are passed by name. return {"squared": squared, "partial": partial} def eval_metrics_fn(partial, squared): return { "mean": metrics_impl.mean(partial), "loss": squared, } return IPUPipelineEstimatorSpec( mode, computational_stages=[stage1, stage2], eval_metrics_fn=eval_metrics_fn, gradient_accumulation_count=gradient_accumulation_count)
def model_fn_with_zero_stages(mode): def optimizer_function(): pass return IPUPipelineEstimatorSpec(mode, computational_stages=[], gradient_accumulation_count=1, optimizer_function=optimizer_function)
def model_fn_without_optimizer_function(mode): def stage1(features, labels): return features, labels def stage2(partial, labels): return partial + labels return IPUPipelineEstimatorSpec(mode, computational_stages=[stage1, stage2], gradient_accumulation_count=1)
def my_model_fn(mode): def stage1(features): w = variable_scope.get_variable("w", initializer=1) partial = w * features return partial def stage2(partial): prediction = partial * partial return prediction, partial return IPUPipelineEstimatorSpec(mode, computational_stages=[stage1, stage2], gradient_accumulation_count=4)
def model_fn_without_loss(mode): def stage1(features): return features * features def stage2(partial): return partial + partial def eval_metrics_fn(prediction): return {"prediction": metrics_impl.mean(prediction)} return IPUPipelineEstimatorSpec(mode, computational_stages=[stage1, stage2], eval_metrics_fn=eval_metrics_fn, gradient_accumulation_count=4)
def my_model_fn(mode): self.assertEqual(model_fn_lib.ModeKeys.PREDICT, mode) def stage1(features): w = variable_scope.get_variable("w", initializer=1) partial = w * features return partial def stage2(partial): prediction = partial * partial return prediction return IPUPipelineEstimatorSpec( mode, computational_stages=[stage1, stage2], gradient_accumulation_count=gradient_accumulation_count)
def my_model_fn(mode): self.assertEqual(model_fn_lib.ModeKeys.TRAIN, mode) def stage1(features, labels): w = variable_scope.get_variable(name="w", initializer=initial_w) partial = w * features return partial, labels def stage2(partial, labels): loss = partial + labels return loss def optimizer_function(loss): opt = gradient_descent.GradientDescentOptimizer(learning_rate) return pipelining_ops.OptimizerFunctionOutput(opt, loss) return IPUPipelineEstimatorSpec( mode, computational_stages=[stage1, stage2], optimizer_function=optimizer_function, gradient_accumulation_count=gradient_accumulation_count)