added runner

models/vgg_rep.py

@@ -75,7 +75,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -101,7 +101,7 @@ def forward(self, x):
         x = self.bn6(x)
         x = F.relu(x)
         x = self.maxpool(x)
-        
+
         x = self.conv4_1(x, self.conv2_1.weight)
         x = self.bn7(x)
         x = F.relu(x)
@@ -148,7 +148,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -221,7 +221,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -294,7 +294,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -367,7 +367,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -440,7 +440,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)
@@ -513,7 +513,7 @@ def __init__(self, num_classes=10, args=None):
         self.bn10 = nn.BatchNorm2d(512)
         self.fc = nn.Linear(512, num_classes)
         self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.avgpool = nn.AvgPool2d(kernel_size=1, stride=1)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
 
     def forward(self, x):
         x = self.conv1_1(x)

runner.py

@@ -0,0 +1,43 @@
+import subprocess as sp
+from multiprocessing.dummy import Pool
+import itertools
+import sys, os
+import argparse
+import random
+
+# adapt these to you setup
+NR_GPUS = 2
+NR_PROCESSES = 2
+
+cnt = -1
+
+
+def call_script(scripts):
+    global cnt
+    model, dataset, num_classes, wa, epochs = scripts
+    crt_env = os.environ.copy()
+    crt_env['OMP_NUM_THREADS'] = '1'
+    crt_env['MKL_NUM_THREADS'] = '1'
+    crt_env['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
+    cnt += 1
+    gpu = cnt % NR_GPUS
+    crt_env['CUDA_VISIBLE_DEVICES'] = str(gpu)
+    print(scripts)
+    sp.call([sys.executable, 'main.py', '--model_name', str(model), '--dataset', str(dataset), '--num_classes', str(num_classes), '--epochs', str(epochs), '--save_weights', '--weight_activation', wa], env=crt_env)
+
+
+if __name__ == '__main__':
+    pool = Pool(NR_PROCESSES)
+    models = [f'rep_vgg_{i}' for i in range(4,11)]
+    was = ['linear', 'bireal', 'swish']
+
+    scripts = list(itertools.product(models, ['TINYIMNET'], [200], was, [90]))
+    models = [f'VGG{i}' for i in range(4,12)]
+    scripts += list(itertools.product(models, ['TINYIMNET'], [200], ['linear'], [90]))
+    scripts += list(itertools.product(models, ['CIFAR100'], [100], ['linear'], [160]))
+    scripts += list(itertools.product(models, ['CIFAR10'], [10], ['linear'], [160]))
+
+    pool.map(call_script, scripts)
+    pool.close()
+    pool.join()
+