Use tqdm() library and remove pytorch_classification Bar (#183)

suragnair · May 6, 2020 · 4a89361 · 4a89361
1 parent f1c98a9
commit 4a89361
Show file tree

Hide file tree

Showing 42 changed files with 121 additions and 3,107 deletions.
diff --git a/Arena.py b/Arena.py
@@ -1,7 +1,6 @@
 import logging
-import time
 
-from pytorch_classification.utils import Bar, AverageMeter
+from tqdm import tqdm
 
 log = logging.getLogger(__name__)
 
@@ -73,56 +72,29 @@ def playGames(self, num, verbose=False):
             twoWon: games won by player2
             draws:  games won by nobody
         """
-        eps_time = AverageMeter()
-        bar = Bar('Arena.playGames', max=num)
-        end = time.time()
-        eps = 0
-        maxeps = int(num)
 
         num = int(num / 2)
         oneWon = 0
         twoWon = 0
         draws = 0
-        for _ in range(num):
+        for _ in tqdm(range(num), desc="Arena.playGames (1)"):
             gameResult = self.playGame(verbose=verbose)
             if gameResult == 1:
                 oneWon += 1
             elif gameResult == -1:
                 twoWon += 1
             else:
                 draws += 1
-            # bookkeeping + plot progress
-            eps += 1
-            eps_time.update(time.time() - end)
-            end = time.time()
-            bar.suffix = '({eps}/{maxeps}) Eps Time: {et:.3f}s | Total: {total:} | ETA: {eta:}'.format(eps=eps,
-                                                                                                       maxeps=maxeps,
-                                                                                                       et=eps_time.avg,
-                                                                                                       total=bar.elapsed_td,
-                                                                                                       eta=bar.eta_td)
-            bar.next()
 
         self.player1, self.player2 = self.player2, self.player1
 
-        for _ in range(num):
+        for _ in tqdm(range(num), desc="Arena.playGames (2)"):
             gameResult = self.playGame(verbose=verbose)
             if gameResult == -1:
                 oneWon += 1
             elif gameResult == 1:
                 twoWon += 1
             else:
                 draws += 1
-            # bookkeeping + plot progress
-            eps += 1
-            eps_time.update(time.time() - end)
-            end = time.time()
-            bar.suffix = '({eps}/{maxeps}) Eps Time: {et:.3f}s | Total: {total:} | ETA: {eta:}'.format(eps=eps,
-                                                                                                       maxeps=maxeps,
-                                                                                                       et=eps_time.avg,
-                                                                                                       total=bar.elapsed_td,
-                                                                                                       eta=bar.eta_td)
-            bar.next()
-
-        bar.finish()
 
         return oneWon, twoWon, draws
diff --git a/Coach.py b/Coach.py
@@ -1,16 +1,15 @@
 import logging
 import os
 import sys
-import time
 from collections import deque
 from pickle import Pickler, Unpickler
 from random import shuffle
 
 import numpy as np
+from tqdm import tqdm
 
 from Arena import Arena
 from MCTS import MCTS
-from pytorch_classification.utils import Bar, AverageMeter
 
 log = logging.getLogger(__name__)
 
@@ -85,23 +84,10 @@ def learn(self):
             if not self.skipFirstSelfPlay or i > 1:
                 iterationTrainExamples = deque([], maxlen=self.args.maxlenOfQueue)
 
-                eps_time = AverageMeter()
-                bar = Bar('Self Play', max=self.args.numEps)
-                end = time.time()
-
-                for eps in range(self.args.numEps):
+                for _ in tqdm(range(self.args.numEps), desc="Self Play"):
                     self.mcts = MCTS(self.game, self.nnet, self.args)  # reset search tree
                     iterationTrainExamples += self.executeEpisode()
 
-                    # bookkeeping + plot progress
-                    eps_time.update(time.time() - end)
-                    end = time.time()
-                    bar.suffix = '({eps}/{maxeps}) Eps Time: {et:.3f}s | Total: {total:} | ETA: {eta:}'.format(
-                        eps=eps + 1, maxeps=self.args.numEps, et=eps_time.avg,
-                        total=bar.elapsed_td, eta=bar.eta_td)
-                    bar.next()
-                bar.finish()
-
                 # save the iteration examples to the history 
                 self.trainExamplesHistory.append(iterationTrainExamples)
 
@@ -151,6 +137,7 @@ def saveTrainExamples(self, iteration):
         filename = os.path.join(folder, self.getCheckpointFile(iteration) + ".examples")
         with open(filename, "wb+") as f:
             Pickler(f).dump(self.trainExamplesHistory)
+        f.closed
 
     def loadTrainExamples(self):
         modelFile = os.path.join(self.args.load_folder_file[0], self.args.load_folder_file[1])

diff --git a/connect4/tensorflow/NNet.py b/connect4/tensorflow/NNet.py
@@ -1,13 +1,12 @@
 import os
-import shutil
+import sys
 import time
-import random
+
 import numpy as np
-import math
-import sys
+from tqdm import tqdm
+
 sys.path.append('../../')
 from utils import *
-from pytorch_classification.utils import Bar, AverageMeter
 from NeuralNet import NeuralNet
 
 import tensorflow as tf
@@ -43,51 +42,26 @@ def train(self, examples):
 
         for epoch in range(args.epochs):
             print('EPOCH ::: ' + str(epoch + 1))
-            data_time = AverageMeter()
-            batch_time = AverageMeter()
             pi_losses = AverageMeter()
             v_losses = AverageMeter()
-            end = time.time()
-
-            bar = Bar('Training Net', max=int(len(examples) / args.batch_size))
-            batch_idx = 0
+            batch_count = int(len(examples) / args.batch_size)
 
             # self.sess.run(tf.local_variables_initializer())
-            while batch_idx < int(len(examples) / args.batch_size):
+            t = tqdm(range(batch_count), desc='Training Net')
+            for _ in t:
                 sample_ids = np.random.randint(len(examples), size=args.batch_size)
                 boards, pis, vs = list(zip(*[examples[i] for i in sample_ids]))
 
                 # predict and compute gradient and do SGD step
-                input_dict = {self.nnet.input_boards: boards, self.nnet.target_pis: pis, self.nnet.target_vs: vs, self.nnet.dropout: args.dropout, self.nnet.isTraining: True}
-
-                # measure data loading time
-                data_time.update(time.time() - end)
+                input_dict = {self.nnet.input_boards: boards, self.nnet.target_pis: pis, self.nnet.target_vs: vs,
+                              self.nnet.dropout: args.dropout, self.nnet.isTraining: True}
 
                 # record loss
                 self.sess.run(self.nnet.train_step, feed_dict=input_dict)
                 pi_loss, v_loss = self.sess.run([self.nnet.loss_pi, self.nnet.loss_v], feed_dict=input_dict)
                 pi_losses.update(pi_loss, len(boards))
                 v_losses.update(v_loss, len(boards))
-
-                # measure elapsed time
-                batch_time.update(time.time() - end)
-                end = time.time()
-                batch_idx += 1
-
-                # plot progress
-                bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss_pi: {lpi:.4f} | Loss_v: {lv:.3f}'.format(
-                            batch=batch_idx,
-                            size=int(len(examples) / args.batch_size),
-                            data=data_time.avg,
-                            bt=batch_time.avg,
-                            total=bar.elapsed_td,
-                            eta=bar.eta_td,
-                            lpi=pi_losses.avg,
-                            lv=v_losses.avg,
-                            )
-                bar.next()
-            bar.finish()
-
+                t.set_postfix(Loss_pi=pi_losses, Loss_v=v_losses)
 
     def predict(self, board):
         """
@@ -100,9 +74,11 @@ def predict(self, board):
         board = board[np.newaxis, :, :]
 
         # run
-        prob, v = self.sess.run([self.nnet.prob, self.nnet.v], feed_dict={self.nnet.input_boards: board, self.nnet.dropout: 0, self.nnet.isTraining: False})
+        prob, v = self.sess.run([self.nnet.prob, self.nnet.v],
+                                feed_dict={self.nnet.input_boards: board, self.nnet.dropout: 0,
+                                           self.nnet.isTraining: False})
 
-        #print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
+        # print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
         return prob[0], v[0]
 
     def save_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
@@ -120,7 +96,7 @@ def save_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
     def load_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
         filepath = os.path.join(folder, filename)
         if not os.path.exists(filepath + '.meta'):
-            raise("No model in path {}".format(filepath))
+            raise ("No model in path {}".format(filepath))
         with self.nnet.graph.as_default():
             self.saver = tf.train.Saver()
             self.saver.restore(self.sess, filepath)
diff --git a/gobang/tensorflow/NNet.py b/gobang/tensorflow/NNet.py
@@ -1,13 +1,11 @@
 import os
-import shutil
-import time
-import random
-import numpy as np
-import math
 import sys
+
+import numpy as np
+from tqdm import tqdm
+
 sys.path.append('../../')
 from utils import *
-from pytorch_classification.utils import Bar, AverageMeter
 from NeuralNet import NeuralNet
 
 import tensorflow as tf
@@ -21,6 +19,7 @@
     'num_channels': 512,
 })
 
+
 class NNetWrapper(NeuralNet):
     def __init__(self, game):
         self.nnet = onnet(game, args)
@@ -39,52 +38,27 @@ def train(self, examples):
         """
 
         for epoch in range(args.epochs):
-            print('EPOCH ::: ' + str(epoch+1))
-            data_time = AverageMeter()
-            batch_time = AverageMeter()
+            print('EPOCH ::: ' + str(epoch + 1))
             pi_losses = AverageMeter()
             v_losses = AverageMeter()
-            end = time.time()
-
-            bar = Bar('Training Net', max=int(len(examples)/args.batch_size))
-            batch_idx = 0
+            batch_count = int(len(examples) / args.batch_size)
 
             # self.sess.run(tf.local_variables_initializer())
-            while batch_idx < int(len(examples)/args.batch_size):
+            t = tqdm(range(batch_count), desc='Training Net')
+            for _ in t:
                 sample_ids = np.random.randint(len(examples), size=args.batch_size)
                 boards, pis, vs = list(zip(*[examples[i] for i in sample_ids]))
 
                 # predict and compute gradient and do SGD step
-                input_dict = {self.nnet.input_boards: boards, self.nnet.target_pis: pis, self.nnet.target_vs: vs, self.nnet.dropout: args.dropout, self.nnet.isTraining: True}
-
-                # measure data loading time
-                data_time.update(time.time() - end)
+                input_dict = {self.nnet.input_boards: boards, self.nnet.target_pis: pis, self.nnet.target_vs: vs,
+                              self.nnet.dropout: args.dropout, self.nnet.isTraining: True}
 
                 # record loss
                 self.sess.run(self.nnet.train_step, feed_dict=input_dict)
                 pi_loss, v_loss = self.sess.run([self.nnet.loss_pi, self.nnet.loss_v], feed_dict=input_dict)
                 pi_losses.update(pi_loss, len(boards))
                 v_losses.update(v_loss, len(boards))
-
-                # measure elapsed time
-                batch_time.update(time.time() - end)
-                end = time.time()
-                batch_idx += 1
-
-                # plot progress
-                bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss_pi: {lpi:.4f} | Loss_v: {lv:.3f}'.format(
-                            batch=batch_idx,
-                            size=int(len(examples)/args.batch_size),
-                            data=data_time.avg,
-                            bt=batch_time.avg,
-                            total=bar.elapsed_td,
-                            eta=bar.eta_td,
-                            lpi=pi_losses.avg,
-                            lv=v_losses.avg,
-                            )
-                bar.next()
-            bar.finish()
-
+                t.set_postfix(Loss_pi=pi_losses, Loss_v=v_losses)
 
     def predict(self, board):
         """
@@ -97,9 +71,11 @@ def predict(self, board):
         board = board[np.newaxis, :, :]
 
         # run
-        prob, v = self.sess.run([self.nnet.prob, self.nnet.v], feed_dict={self.nnet.input_boards: board, self.nnet.dropout: 0, self.nnet.isTraining: False})
+        prob, v = self.sess.run([self.nnet.prob, self.nnet.v],
+                                feed_dict={self.nnet.input_boards: board, self.nnet.dropout: 0,
+                                           self.nnet.isTraining: False})
 
-        #print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
+        # print('PREDICTION TIME TAKEN : {0:03f}'.format(time.time()-start))
         return prob[0], v[0]
 
     def save_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
@@ -109,15 +85,15 @@ def save_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
             os.mkdir(folder)
         else:
             print("Checkpoint Directory exists! ")
-        if self.saver == None:            
+        if self.saver == None:
             self.saver = tf.train.Saver(self.nnet.graph.get_collection('variables'))
         with self.nnet.graph.as_default():
             self.saver.save(self.sess, filepath)
 
     def load_checkpoint(self, folder='checkpoint', filename='checkpoint.pth.tar'):
         filepath = os.path.join(folder, filename)
-        if not os.path.exists(filepath+'.meta'):
-            raise("No model in path {}".format(filepath))
+        if not os.path.exists(filepath + '.meta'):
+            raise ("No model in path {}".format(filepath))
         with self.nnet.graph.as_default():
             self.saver = tf.train.Saver()
-            self.saver.restore(self.sess, filepath)
+            self.saver.restore(self.sess, filepath)