sizer.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import os, stat, sys
   9 import time
  10 import torch
  11 from torch import nn
  12
  13 t = 0
  14
  15 if len(sys.argv) < 2:
  16     print(sys.argv[0] + ''' <file to monitor>
  17
  18 For example:
  19
  20 (17, 3, 60, 80)
  21 nn.Conv2d(3, 32, 3, padding = 1)
  22 nn.MaxPool2d(2)
  23 nn.Conv2d(32, 32, 3, padding = 1)
  24 nn.MaxPool2d(2)
  25 nn.Conv2d(32, 64, 3, padding = 1)
  26 nn.MaxPool2d(5)
  27 nn.Conv2d(64, 64, (3, 4))''')
  28     exit(1)
  29
  30 while True:
  31     pt = t
  32     t = os.stat(sys.argv[1])[stat.ST_MTIME]
  33     if t > pt:
  34         pt = t
  35         os.system('clear')
  36         try:
  37             temp = [l.strip('\n\r') for l in open(sys.argv[1], 'r').readlines()]
  38             x = torch.zeros(eval(temp.pop(0)))
  39             print('-> ' + str(tuple(x.size())))
  40             for k in temp:
  41                 print('   ' + k)
  42                 x = eval(k + '(x)')
  43                 print('-> ' + str(tuple(x.size())))
  44         except:
  45             print('** Error **')
  46     time.sleep(1)