|
| 1 | +using NumSharp; |
| 2 | +using System; |
| 3 | +using System.Collections.Generic; |
| 4 | +using System.Linq; |
| 5 | +using System.Text; |
| 6 | +using Tensorflow; |
| 7 | +using Tensorflow.Graphs; |
| 8 | +using Tensorflow.Keras.ArgsDefinition; |
| 9 | +using Tensorflow.Keras.Engine; |
| 10 | +using static Tensorflow.Binding; |
| 11 | + |
| 12 | +namespace Tensorflow.Keras.Layers |
| 13 | +{ |
| 14 | + public class TensorFlowOpLayer : Layer |
| 15 | + { |
| 16 | + TensorFlowOpLayerArgs args; |
| 17 | + Dictionary<int, NDArray> constants => args.Constants; |
| 18 | + NodeDef node_def => args.NodeDef; |
| 19 | + static string TF_OP_LAYER_NAME_PREFIX = "tf_op_layer_"; |
| 20 | + public string OpType => node_def.Op; |
| 21 | + |
| 22 | + public TensorFlowOpLayer(TensorFlowOpLayerArgs args) |
| 23 | + : base(new LayerArgs |
| 24 | + { |
| 25 | + Name = TF_OP_LAYER_NAME_PREFIX + args.Name, |
| 26 | + Trainable = args.Trainable, |
| 27 | + DType = args.DType, |
| 28 | + Autocast = false |
| 29 | + }) |
| 30 | + { |
| 31 | + this.args = args; |
| 32 | + built = true; |
| 33 | + } |
| 34 | + |
| 35 | + protected override Tensors Call(Tensors inputs, Tensor state = null, bool is_training = false) |
| 36 | + { |
| 37 | + if (tf.Context.executing_eagerly()) |
| 38 | + return _defun_call(inputs); |
| 39 | + return MakOp(inputs); |
| 40 | + } |
| 41 | + |
| 42 | + [AutoGraph] |
| 43 | + Tensors _defun_call(Tensors inputs) |
| 44 | + => MakOp(inputs); |
| 45 | + |
| 46 | + Tensors MakOp(Tensors inputs) |
| 47 | + { |
| 48 | + var graph = inputs.graph; |
| 49 | + graph.as_default(); |
| 50 | + foreach (var (index, constant) in enumerate(constants)) |
| 51 | + { |
| 52 | + var value = constant_op.constant(constant, name: node_def.Input[index]); |
| 53 | + inputs.Insert(index, value); |
| 54 | + } |
| 55 | + |
| 56 | + var (c_op, _) = ops._create_c_op(graph, node_def, inputs.ToArray(), new Operation[0]); |
| 57 | + var op = graph._create_op_from_tf_operation(c_op); |
| 58 | + op._control_flow_post_processing(); |
| 59 | + |
| 60 | + // Record the gradient because custom-made ops don't go through the |
| 61 | + // code-gen'd eager call path |
| 62 | + var op_type = op.node_def.Op; |
| 63 | + |
| 64 | + tf.Runner.RecordGradient(op_type, op.inputs._inputs, null, op.outputs); |
| 65 | + |
| 66 | + graph.Exit(); |
| 67 | + return op.outputs; |
| 68 | + } |
| 69 | + |
| 70 | + public Layer GetOpLayer(TensorFlowOpLayerArgs args) |
| 71 | + => new TensorFlowOpLayer(args); |
| 72 | + } |
| 73 | +} |
0 commit comments