periodically savingin out model and adding functionality to make identity Phi networks (for 0-dim NNs)

src/bc/bc.py

@@ -22,7 +22,7 @@ class BehaviorCloningPolicy():
         """
         self._config = config
         self._transforms = transforms
-        self._policy = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
+        self._policy = DeepSetsPolicy(config)
         
     @property
     def transforms(self):
@@ -85,13 +85,15 @@ class BehaviorCloningPolicy():
     def parameters(self):
         return self._policy.parameters()
 
-    def save_model(self, filestr):
+    def save_model(self, filestr, save_transforms=True):
         """
         Save transforms and state_dict to a location specificed by filestr
         Args:
             filestr (str): string prefix to save model to
+            save_transforms (bool): whether to save transforms
         """
-        pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
+        if save_transforms:
+            pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
         torch.save(self._policy.state_dict(), filestr+'_model.pt')
 
 def generate_transforms(dataset):
@@ -116,13 +118,16 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
     
     # hyperparams
     train_epochs = 1000
-    cv_every = 10
-    epoch_every = 1000
     train_batch_size = 64
-    cv_batch_size = 256 # doesn't matter
     learning_rate = 1e-3
     weight_decay = 0.1
 
+    cv_every = 10
+    print_epoch_every = 1000
+    print_cv_every = 1000
+    checkpoint_every = 100
+    cv_batch_size = 256 # doesn't matter
+
     # generate transform from train_dataset
     transforms = generate_transforms(train_dataset)
 
@@ -165,7 +170,7 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
         # Write epoch loss
         
         writer.add_scalar('training loss',epoch_loss, i)
-        if i % epoch_every == 0:
+        if i % print_epoch_every == 0:
             print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
 
         # measure cv loss
@@ -177,6 +182,10 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
                     loss = loss_fn(pred_action, batch['action'])
                     cv_loss += loss.item() / len(cv_dataset)
             writer.add_scalar('cv loss', cv_loss, i)
+        if i % print_cv_every == 0:    
             print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
 
+        # save model checkpoints
+        if i % checkpoint_every == 0:
+            policy.save_model(filestr + '_epoch%04i'%(i) )
     policy.save_model(filestr)

src/nets/deepsets.py

@@ -90,10 +90,14 @@ class Phi(nn.Module):
         super(Phi, self).__init__()
         self.input_dim = input_dim
         self.output_dim = output_dim
-        self.layers = nn.ModuleList([nn.Linear(self.input_dim, hidden_dim)])
-        for _ in range(hidden_n - 1):
-            self.layers.append(nn.Linear(hidden_dim, hidden_dim))
-        self.layers.append(nn.Linear(hidden_dim, self.output_dim))
+        if hidden_n > 0:
+            self.layers = nn.ModuleList([nn.Linear(self.input_dim, hidden_dim)])
+            for _ in range(hidden_n - 1):
+                self.layers.append(nn.Linear(hidden_dim, hidden_dim))
+            self.layers.append(nn.Linear(hidden_dim, self.output_dim))
+        else:
+            self.layers = nn.ModuleList([nn.Identity()])
+            self.output_dim = self.input_dim
         self.activation = nn.functional.relu
         self.final_activation = final_activation if final_activation else lambda x: x
 

src/policies/policy.py

@@ -8,20 +8,23 @@ class Policy:
     pass
 
 class DeepSetsPolicy(Policy, nn.Module):
-    def __init__(self, ego_config, deepsets_config, path_config, head_config):
+    def __init__(self, config):
         """
         Args:
-            ego_config (dict): dictionary for configuring the ego network
-            deepsets_config (dict): dictionary for configuring the deepsets network
-            path_config (dict): dictionary for configuring the path network
-            head_config (dict): dictionary for configuring the common head network
+            config (dict): dictionary for configuring the deep sets policy
         """
         super(DeepSetsPolicy, self).__init__()
-        self.ego_net = Phi.from_config(ego_config)
-        self.deepsets_net = DeepSetsModule.from_config(deepsets_config)
-        self.path_net = Phi.from_config(path_config)
+        ego_config = config['ego_state']
+        deepsets_config = config['deepsets']
+        pathnet_config = config['path_encoder']
+
+        self.ego_net = Phi.from_config(ego_config) if ego_config else lambda x: x
+        self.deepsets_net = DeepSetsModule.from_config(deepsets_config) if deepsets_config else lambda x: x
+        self.path_net = Phi.from_config(pathnet_config) if pathnet_config else lambda x: x
+
         cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim
         # head has number of concatenated features as input
+        head_config = config['head']
         head_config["input_dim"] = cat_dim
         self.head = Phi.from_config(head_config)