Added dynamic dumping and loading of trainExamples for memory optimiz…

…ation

GoCoach.py

@@ -20,7 +20,7 @@ class Coach():
     in Game and NeuralNet. args are specified in main.py.
     """
 
-    def __init__(self, game, nnet, args, log=False, logPath=''):
+    def __init__(self, game, nnet, args, log=False, logPath='', NetType='RES'):
         self.game = game
         self.nnet = nnet
         self.pnet = self.nnet.__class__(self.game, t=self.nnet.netType)  # the competitor network
@@ -36,6 +36,7 @@ def __init__(self, game, nnet, args, log=False, logPath=''):
         self.winRate = []
         self.currentEpisode = 0
         self.iterationTrainExamples = deque([], maxlen=self.args.maxlenOfQueue)
+        self.NetType = NetType
 
     def executeEpisode(self):
         """
@@ -186,17 +187,33 @@ def learn(self):
             counts_file.write(f"\n Number of games added to train examples during iteration #{i}: {downloads_count} games\n")
             counts_file.close()
 
-            # save the iteration examples to the history
-            if not self.skipFirstSelfPlay or append_downloads:
-                self.trainExamplesHistory.append(self.iterationTrainExamples)
+            if i != 1 or self.args.load_model:
+                new_train_examples = self.iterationTrainExamples
 
-            if len(self.trainExamplesHistory) > self.args.numItersForTrainExamplesHistory:
+                # update pathing
+                checkpoint_dir = f'logs/go/{self.NetType}_MCTS_SimModified_checkpoint/{self.game.getBoardSize()[0]}/'
+                checkpoint_files = [file for file in os.listdir(checkpoint_dir) if
+                                    file.startswith('checkpoint_') and file.endswith('.pth.tar.examples')]
+                latest_checkpoint = max(checkpoint_files, key=lambda x: int(x.split('_')[1].split('.')[0]))
+                self.args.load_folder_file = [
+                    f'logs/go/{self.NetType}_MCTS_SimModified_checkpoint/{self.game.getBoardSize()[0]}/',
+                    latest_checkpoint]
+
+                self.loadTrainExamples()
+
+                self.trainExamplesHistory.append(new_train_examples)
+            else:
+                # save the iteration examples to the history
+                if not self.skipFirstSelfPlay or append_downloads:
+                    self.trainExamplesHistory.append(self.iterationTrainExamples)
+
+            while len(self.trainExamplesHistory) > self.args.numItersForTrainExamplesHistory:
                 # print("len(trainExamplesHistory) =", len(self.trainExamplesHistory), " => remove the oldest trainExamples")
                 self.trainExamplesHistory.pop(0)
 
             # backup history to a file
             # NB! the examples were collected using the model from the previous iteration, so (i-1)
-            self.saveTrainExamples(i - 1)
+            self.saveTrainExamples(i-1)
 
             # shuffle examples before training
             trainExamples = []
@@ -210,12 +227,18 @@ def learn(self):
             pmcts = MCTS(self.game, self.pnet, self.args)
 
             trainLog = self.nnet.train(trainExamples)
+
+            # clear trainExamples to save memory, reload when needed
+            self.trainExamplesHistory = []
+            self.iterationTrainExamples.clear()
+
             self.p_loss_per_iteration.append(np.average(trainLog['P_LOSS'].to_numpy()))
             self.v_loss_per_iteration.append(np.average(trainLog['V_LOSS'].to_numpy()))
             if self.keepLog:
                 trainLog.to_csv(self.logPath + 'ITER_{}_TRAIN_LOG.csv'.format(i))
 
             iterHistory['ITER_DETAIL'].append(self.logPath + 'ITER_{}_TRAIN_LOG.csv'.format(i))
+
             nmcts = MCTS(self.game, self.nnet, self.args)
 
             print('\nPITTING AGAINST PREVIOUS VERSION')
@@ -322,7 +345,7 @@ def loadLosses(self):
             if r != "y":
                 sys.exit()
         else:
-            print("File with trainExamples found. Read it.")
+            print("File with graph information found. Read it.")
             with open(vlossFile, "rb") as f:
                 self.v_loss_per_iteration = Unpickler(f).load()
             f.closed

main.py

@@ -23,19 +23,19 @@ class Display(IntEnum):
 args = dotdict({
     # training parameters
     'numIters': 1000,
-    'numEps': 2,  # Number of complete self-play games to simulate during a new iteration.
+    'numEps': 20,  # Number of complete self-play games to simulate during a new iteration.
     'tempThreshold': 15,
     'updateThreshold': 0.54,
     # During arena playoff, new neural net will be accepted if threshold or more of games are won.
     'maxlenOfQueue': 200000,  # Number of game examples to train the neural networks.
     'numMCTSSims': 150,  # Number of games moves for MCTS to simulate.
     'arenaCompare': 2,  # Number of games to play during arena play to determine if new net will be accepted.
     'cpuct': 1.0,
-    'numItersForTrainExamplesHistory': 25,
+    'numItersForTrainExamplesHistory': 15,
 
     # customization
-    'load_model': False,
-    'distributed_training': False,  # use Google Drivr for computing on multiple machines
+    'load_model': True,
+    'distributed_training': False,  # use Google Drive for computing on multiple machines
     'display': Display.DISPLAY_BAR,
 
     # utility
@@ -46,13 +46,15 @@ class Display(IntEnum):
     'start_time': time.time(),
     'checkpoint': './logs/go/{}_checkpoint/{}/'.format(NetType + '_' + tag, BoardSize),
 })
+
 if args.load_model:
     checkpoint_dir = f'logs/go/{NetType}_MCTS_SimModified_checkpoint/{BoardSize}/'
     checkpoint_files = [file for file in os.listdir(checkpoint_dir) if
                         file.startswith('checkpoint_') and file.endswith('.pth.tar.examples')]
     latest_checkpoint = max(checkpoint_files, key=lambda x: int(x.split('_')[1].split('.')[0]))
     args['load_folder_file'] = [f'logs/go/{NetType}_MCTS_SimModified_checkpoint/{BoardSize}/', latest_checkpoint]
     args['start_iter'] = int(latest_checkpoint.split('_')[1].split('.')[0]) + 1
+
 else:
     args['start_iter'] = 1
 
@@ -89,10 +91,9 @@ class Display(IntEnum):
         # instead of the load_checkpoint() function
         nnet.load_checkpoint(args.checkpoint, 'best.pth.tar')
 
-    c = Coach(g, nnet, args, log=True, logPath=logPath)
+    c = Coach(g, nnet, args, log=True, logPath=logPath, NetType=NetType)
 
     if args.load_model:
-        print("Loading trainExamples from file")
-        c.loadTrainExamples()
+        c.skipFirstSelfPlay = True
 
     c.learn()