From 69359b5af36e6512e480930d1b31268e12b9ba77 Mon Sep 17 00:00:00 2001 From: Arec Date: Mon, 26 Jul 2021 05:51:08 -0700 Subject: [PATCH] periodically savingin out model and adding functionality to make identity Phi networks (for 0-dim NNs) --- src/bc/bc.py | 23 ++++++++++++++++------- src/nets/deepsets.py | 12 ++++++++---- src/policies/policy.py | 19 +++++++++++-------- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/src/bc/bc.py b/src/bc/bc.py index b299c86..0fa58dd 100644 --- a/src/bc/bc.py +++ b/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) diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py index 82ba7a8..d16ff6d 100644 --- a/src/nets/deepsets.py +++ b/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 diff --git a/src/policies/policy.py b/src/policies/policy.py index 2ac0f22..dad8f92 100644 --- a/src/policies/policy.py +++ b/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)