Refactor n frames

gferrate · gferrate · commit b20080f849f9 · 2021-01-21T19:01:25.000+01:00
diff --git a/README.md b/README.md
@@ -39,6 +39,7 @@ You can see the results of the experiments done [here](results/README.md).
         - There are 3 env variables which can be set:
             - `DELETE_OLD_LOGS`: if set to 1, removes all logs in `logs/` folder. Useful for new run.
             - `RESTART_SCREEN`: if set to 1, kills all current screens. Useful for new run.
+            - `N_FRAMES`: Number of input frames.
             - `SPLIT_TYPE`: Can be set to 3 values (more information in dissemination):
                 - `no_split`: No split in the dataset is performed.
                 - `iid`: IID split type is performed.
diff --git a/classification/trainer.py b/classification/trainer.py
@@ -9,7 +9,6 @@
 
 
 epochs = 1  # Before 30
-n_frames = 7  # Optimal == 7
 batch_size = 10
 workers = 0
 
@@ -19,9 +18,31 @@
 
 
 class Trainer(object):
-    def __init__(self):
-        # __init__ overrided in child classes
-        pass
+
+    def __init__(self, port, n_frames, client_id, num_clients, data_split_type):
+        # Common init in both childs
+        self.port = port
+        self.n_frames = n_frames  # Optimal frames == 7
+        self.client_id = client_id
+        self.num_clients = num_clients
+        self.data_split_type = data_split_type
+        self.logger = 'logs/{}.log'.format(self.client_id)
+        self.init_logger()
+        self.metaFile = self.get_meta_file()
+
+    def get_meta_file(self):
+        mf = 'data/classification'
+        if self.data_split_type == 'iid':
+            return '{}/metadata_{}_clients_iid.mat'.format(mf,
+                                                           self.num_clients)
+        elif self.data_split_type == 'non-iid-a':
+            return '{}/metadata_{}_clients_non_iid_a.mat'.format(
+                mf, self.num_clients)
+        elif self.data_split_type == 'no_split':
+            return '{}/metadata.mat'.format(mf)
+        else:
+            raise Exception('Data split type "{}" not implemented'.format(
+                self.data_split_type))
 
     def init_logger(self):
         logging.basicConfig(
@@ -40,7 +61,7 @@ def loadDatasets(self,
         return torch.utils.data.DataLoader(
             ObjectClusterDataset(
                 split=split, doAugment=(split == 'train'),
-                doFilter=doFilter, sequenceLength=n_frames,
+                doFilter=doFilter, sequenceLength=self.n_frames,
                 metaFile=self.metaFile, useClusters=useClusterSampling
             ),
             batch_size=batch_size,
@@ -114,7 +135,7 @@ def initModel(self):
 
         self.model = Model(
             numClasses=len(self.val_loader.dataset.meta['objects']),
-            sequenceLength=n_frames)
+            sequenceLength=self.n_frames)
         self.model.epoch = 0
         self.model.bestPrec = -1e20
 
@@ -213,28 +234,21 @@ def save_model(self):
 
 
 class SecAggTrainer(Trainer):
-    def __init__(self, client_id, num_clients, data_split_type):
-        self.client_id = client_id
-        self.logger = 'logs/{}.log'.format(self.client_id)
-        self.init_logger()
+    def __init__(self,
+                 port,
+                 n_frames,
+                 num_clients,
+                 data_split_type):
+        client_id = 'sec_agg'
+        super().__init__(port,
+                         n_frames,
+                         client_id,
+                         num_clients,
+                         data_split_type)
         self.type = 'secure_aggregator'
         self.snapshotDir = 'secure_aggregator/persistent_storage'
-        self.client_number = None
         self.train_split = 'train' # Shouldn't be needed since it doesn't train
-        mf = 'data/classification'
-        if data_split_type == 'iid':
-            self.metaFile = '{}/metadata_{}_clients_iid.mat'.format(
-                mf, num_clients)
-        elif data_split_type == 'non-iid-a':
-            self.metaFile = '{}/metadata_{}_clients_non_iid_a.mat'.format(
-                mf, num_clients)
-        elif data_split_type == 'no_split':
-            self.metaFile = '{}/metadata.mat'.format(mf)
-        else:
-            raise Exception('Data split type "{}" not implemented'.format(
-                data_split_type))
         self.init()
-        super(Trainer, self).__init__()
 
     def get_checkpoint_path(self):
         return os.path.join(self.snapshotDir, 'checkpoint.tar')
@@ -244,31 +258,25 @@ def get_best_model_path(self):
 
 
 class ClientTrainer(Trainer):
-    def __init__(self, client_number, client_id, num_clients, data_split_type):
-        self.client_id = client_id
-        self.logger = 'logs/{}.log'.format(self.client_id)
-        self.init_logger()
+    def __init__(self,
+                 port,
+                 n_frames,
+                 client_number,
+                 num_clients,
+                 data_split_type,
+                 client_id):
+        super().__init__(port,
+                         n_frames,
+                         client_id,
+                         num_clients,
+                         data_split_type)
         self.type = 'client'
-        self.snapshotDir = 'client/snapshots_{}'.format(self.client_id)
-        self.client_number = client_number
-        mf = 'data/classification'
-        if data_split_type == 'iid':
-            self.train_split = 'train_{}'.format(client_number)
-            self.metaFile = '{}/metadata_{}_clients_iid.mat'.format(
-                mf, num_clients)
-        elif data_split_type == 'non-iid-a':
-            self.train_split = 'train_{}'.format(client_number)
-            self.metaFile = '{}/metadata_{}_clients_non_iid_a.mat'.format(
-                mf, num_clients)
-        elif data_split_type == 'no_split':
-            self.train_split = 'train'
-            self.metaFile = '{}/metadata.mat'.format(mf)
-        else:
-            raise Exception('Data split type "{}" not implemented'.format(
-                data_split_type))
+        self.snapshotDir = 'client/snapshots_{}'.format(self.port)
+        # Dont do it with the client_id to avoid tons of folders generated
+        self.train_split = self.get_train_split()
 
         # Split dataset if file does not exist
-        if data_split_type in ('iid', 'non-iid-a', 'non-iid-b'):
+        if self.data_split_type in ('iid', 'non-iid-a', 'non-iid-b'):
             from shared import dataset_tools
             # TODO: Reimplement this with a lock file.
             # If multiple clients are spawned this can be a problem.
@@ -287,7 +295,14 @@ def __init__(self, client_number, client_id, num_clients, data_split_type):
                 logging.info('File {} already exists. '
                              'Not creating.'.format(self.metaFile))
         self.init()
-        super(Trainer, self).__init__()
+
+    def get_train_split(self):
+        if self.data_split_type == 'iid':
+            return 'train_{}'.format(self.client_number)
+        elif self.data_split_type == 'non-iid-a':
+            return 'train_{}'.format(self.client_number)
+        elif self.data_split_type == 'no_split':
+            return 'train'
 
 
 class AverageMeter(object):
diff --git a/client/app.py b/client/app.py
@@ -25,6 +25,9 @@
                     default='no_split',
                     help=('Metadata split type. '
                           'Example: no_split, iid, non-iid-a'))
+parser.add_argument('-f', '--n-frames', type=int, required=False,
+                    default=1,
+                    help='n_frames. [1-8]. Default=1')
 
 rsa = rsa_utils.RSAUtils()
 args = parser.parse_args()
@@ -45,8 +48,9 @@
 # _id auto generated by State
 state = State('client', port, _id=None)
 
-client = Client(args.client_number,
-                port,
+client = Client(port,
+                args.n_frames,
+                args.client_number,
                 num_clients,
                 args.split_type,
                 state._id)
diff --git a/initialize.sh b/initialize.sh
@@ -27,21 +27,31 @@ else
     TO_APPEND="-s no_split"
 fi
 
+# N frames
+if [[ -z "$N_FRAMES" ]]
+then
+    echo "No N_FRAMES DEFINED. DEFAULT = 1"
+    N_FRAMES_APPEND=""
+else
+    echo "N_FRAMES DEFINED: $N_FRAMES"
+    N_FRAMES_APPEND=" -f $N_FRAMES"
+fi
+
 # Restart the frontend
 curl http://95.179.192.253:8002/restart
 
 # Start main server in a new screen.
 screen -dmS main_server bash -c "$PYTHON_PATH main_server/app.py -p 8000"
 
 # Start secure aggregator in a new screen.
-screen -dmS secure_aggregator bash -c "$PYTHON_PATH secure_aggregator/app.py -p 8001 $TO_APPEND"
+screen -dmS secure_aggregator bash -c "$PYTHON_PATH secure_aggregator/app.py -p 8001 $TO_APPEND $N_FRAMES_APPEND"
 
 # Start N clients in new screens. Add or comment lines as wanted.
-screen -dmS client_0 bash -c "$PYTHON_PATH client/app.py -p 8003 -n 0 $TO_APPEND"
-screen -dmS client_1 bash -c "$PYTHON_PATH client/app.py -p 8004 -n 1 $TO_APPEND"
-screen -dmS client_2 bash -c "$PYTHON_PATH client/app.py -p 8005 -n 2 $TO_APPEND"
-screen -dmS client_3 bash -c "$PYTHON_PATH client/app.py -p 8006 -n 3 $TO_APPEND"
-screen -dmS client_4 bash -c "$PYTHON_PATH client/app.py -p 8007 -n 4 $TO_APPEND"
+screen -dmS client_0 bash -c "$PYTHON_PATH client/app.py -p 8003 -n 0 $TO_APPEND $N_FRAMES_APPEND"
+screen -dmS client_1 bash -c "$PYTHON_PATH client/app.py -p 8004 -n 1 $TO_APPEND $N_FRAMES_APPEND"
+screen -dmS client_2 bash -c "$PYTHON_PATH client/app.py -p 8005 -n 2 $TO_APPEND $N_FRAMES_APPEND"
+screen -dmS client_3 bash -c "$PYTHON_PATH client/app.py -p 8006 -n 3 $TO_APPEND $N_FRAMES_APPEND"
+screen -dmS client_4 bash -c "$PYTHON_PATH client/app.py -p 8007 -n 4 $TO_APPEND $N_FRAMES_APPEND"
 
 # Start the orchestrator (will start the training).
 echo "Waiting 1 minutes so all clients start"
diff --git a/secure_aggregator/app.py b/secure_aggregator/app.py
@@ -23,13 +23,16 @@
                     default='no_split',
                     help=('Metadata split type. '
                           'Example: no_split, iid, non-iid-a'))
+parser.add_argument('-f', '--n-frames', type=int, required=False,
+                    default=1,
+                    help='n_frames. [1-8]. Default=1')
 
 rsa = rsa_utils.RSAUtils()
 args = parser.parse_args()
 hosts = utils.read_hosts()
 num_clients = len(hosts['clients'])
 
-sec_agg = SecAgg(args.port, num_clients, args.split_type)
+sec_agg = SecAgg(args.port, args.n_frames, num_clients, args.split_type)
 
 # _id auto generated by State
 state = State('secure_aggregator', args.port, _id=None)
