GPT-2: enable passing custom session config for training

Author: lostmsu

GPT-2/Gpt2Trainer.cs

@@ -47,8 +47,11 @@ public Gpt2Trainer(DataSet dataset, Gpt2Encoder encoder, HParams hParams,
         public int SampleEvery { get; set; } = 100;
         public int SampleNum { get; set; } = 1;
 
-        public void Train(string checkpoint, string run, int? counter, CancellationToken cancellation) {
-            new Session().UseSelf(session => {
+        public void Train(string checkpoint, string run, int? counter, dynamic sessionConfig = null, CancellationToken cancellation = default) {
+            Session sess = sessionConfig == null
+                ? Session.NewDyn(config: sessionConfig)
+                : new Session();
+            sess.UseSelf(session => {
                 var context = tf.placeholder(tf.int32, new TensorShape(this.batchSize, null));
                 var output = Gpt2Model.Model(this.hParams, input: context);
                 Tensor labels = context[Range.All, Range.StartAt(1)];

GPT-2/TrainCommand.cs

@@ -8,6 +8,7 @@
     using ManyConsole.CommandLineUtils;
     using numpy;
     using Python.Runtime;
+    using tensorflow.core.protobuf.config_pb2;
     using DataSet = System.Collections.Generic.List<numpy.ndarray>;
     class TrainCommand: ConsoleCommand {
         public override int Run(string[] remainingArguments) {
@@ -34,12 +35,15 @@ public override int Run(string[] remainingArguments) {
             var random = this.Seed == null ? new Random() : new Random(this.Seed.Value);
             var stop = new CancellationTokenSource();
             Console.CancelKeyPress += delegate { stop.Cancel(); };
+            dynamic config = config_pb2.ConfigProto();
+            config.gpu_options.allow_growth = true;
             new Gpt2Trainer(dataset, encoder, hParams, this.BatchSize, this.SampleLength, random) {
                 SaveEvery = this.SaveEvery,
                 SampleNum = this.SampleNum,
                 SampleEvery = this.SampleEvery,
             }
                 .Train(checkpoint, this.RunName,
+                    sessionConfig: config,
                     counter: checkpoint == "fresh" ? 1 : (int?)null,
                     cancellation: stop.Token);