Added option for SELU activation

Author: oscarknagg

experiments/maml.py

@@ -31,6 +31,7 @@
 parser.add_argument('--meta-lr', default=0.005, type=float)
 parser.add_argument('--meta-batch-size', default=32, type=int)
 parser.add_argument('--order', default=1, type=int)
+parser.add_argument('--activation', default='relu', type=str)
 
 args = parser.parse_args()
 
@@ -52,7 +53,7 @@
 meta_batches_per_epoch = 100
 
 param_str = f'{args.dataset}_order={args.order}_n={args.n}_k={args.k}_metabatch={args.meta_batch_size}_' \
-            f'train_steps={args.inner_train_steps}_val_steps={args.inner_val_steps}'
+            f'train_steps={args.inner_train_steps}_val_steps={args.inner_val_steps}_act={args.activation}'
 print(param_str)
 
 
@@ -112,7 +113,6 @@ def prepare_meta_batch_(batch):
         # MAML kwargs
         inner_train_steps=args.inner_val_steps,
         inner_lr=args.inner_lr,
-        num_input_channels=num_input_channels,
         device=device,
         order=args.order,
     ),

few_shot/datasets.py

@@ -13,6 +13,11 @@
 
 class OmniglotDataset(Dataset):
     def __init__(self, subset):
+        """Dataset class representing Omniglot dataset
+
+        # Arguments:
+            subset: Whether the dataset represents the background or evaluation set
+        """
         if subset not in ('background', 'evaluation'):
             raise(ValueError, 'subset must be one of (background, evaluation)')
         self.subset = subset
@@ -49,53 +54,16 @@ def __len__(self):
     def num_classes(self):
         return len(self.df['class_name'].unique())
 
-    def build_n_shot_task(self, k, n=1, query=1):
-        """
-        This method builds a k-way n-shot classification task. It returns a support set of n audio samples each from k
-        unique speakers. In addition it will return a query sample. Downstream models will attempt to match the query
-        sample to the correct samples in the support set.
-        :param k: Number of unique speakers to include in this task
-        :param n: Number of audio samples to include from each speaker
-        :param query: Number of query samples
-        :return:
-        """
-        if k >= self.num_classes():
-            raise(ValueError, 'k must be smaller than the number of unique speakers in this dataset!')
-
-        if k <= 1:
-            raise(ValueError, 'k must be greater than or equal to one!')
-
-        query = self.df.sample(query)
-        query_samples = self[query['id'].values[0]]
-        # Add batch dimension
-        query_samples = (query_samples[0][np.newaxis, :, :], query_samples[1])
-
-        is_query_character = self.df['class_id'] == query['class_id'].values[0]
-        not_query_sample = ~self.df.index.isin(query['id'].values)
-        correct_samples = self.df[is_query_character & not_query_sample].sample(n)
-
-        # Sample k-1 speakers
-        other_support_set_characters = np.random.choice(
-            self.df[~is_query_character]['class_id'].unique(), k-1, replace=False)
-
-        other_support_samples = []
-        for i in range(k-1):
-            is_same_speaker = self.df['class_id'] == other_support_set_characters[i]
-            other_support_samples.append(
-                self.df[~is_query_character & is_same_speaker].sample(n)
-            )
-        support_set = pd.concat([correct_samples]+other_support_samples)
-        support_set_samples = tuple(np.stack(i) for i in zip(*[self[i] for i in support_set.index]))
-
-        return query_samples, support_set_samples
-
     @staticmethod
     def index_subset(subset):
-        """
-        Index a subset by looping through all of it's files and recording their speaker ID, filepath and length.
-        :param subset: Name of the subset
-        :return: A list of dicts containing information about all the audio files in a particular subset of the
-        LibriSpeech dataset
+        """Index a subset by looping through all of its files and recording relevant information.
+
+        # Arguments
+            subset: Name of the subset
+
+        # Returns
+            A list of dicts containing information about all the image files in a particular subset of the
+            Omniglot dataset dataset
         """
         images = []
         print('Indexing {}...'.format(subset))
@@ -127,6 +95,11 @@ def index_subset(subset):
 
 class MiniImageNet(Dataset):
     def __init__(self, subset):
+        """Dataset class representing miniImageNet dataset
+
+        # Arguments:
+            subset: Whether the dataset represents the background or evaluation set
+        """
         if subset not in ('background', 'evaluation'):
             raise(ValueError, 'subset must be one of (background, evaluation)')
         self.subset = subset
@@ -168,6 +141,15 @@ def num_classes(self):
 
     @staticmethod
     def index_subset(subset):
+        """Index a subset by looping through all of its files and recording relevant information.
+
+        # Arguments
+            subset: Name of the subset
+
+        # Returns
+            A list of dicts containing information about all the image files in a particular subset of the
+            miniImageNet dataset
+        """
         images = []
         print('Indexing {}...'.format(subset))
         # Quick first pass to find total for tqdm bar
@@ -196,6 +178,16 @@ def index_subset(subset):
 
 class DummyDataset(Dataset):
     def __init__(self, samples_per_class=10, n_classes=10, n_features=1):
+        """Dummy dataset for debugging/testing purposes
+
+        A sample from the DummyDataset has (n_features + 1) features. The first feature is the index of the sample
+        in the data and the remaining features are the class index.
+
+        # Arguments
+            samples_per_class: Number of samples per class in the dataset
+            n_classes: Number of distinct classes in the dataset
+            n_features: Number of extra features each sample should have.
+        """
         self.samples_per_class = samples_per_class
         self.n_classes = n_classes
         self.n_features = n_features

few_shot/models.py

@@ -37,67 +37,83 @@ def forward(self, input):
         return nn.functional.avg_pool2d(input, kernel_size=input.size()[2:]).view(-1, input.size(1))
 
 
-def conv_block(in_channels, out_channels):
+def conv_block(in_channels, out_channels, activation='relu'):
     """Returns a Module that performs 3x3 convolution, ReLu activation, 2x2 max pooling.
 
     # Arguments
         in_channels:
         out_channels
     """
-    return nn.Sequential(
-        nn.Conv2d(in_channels, out_channels, 3, padding=1),
-        nn.BatchNorm2d(out_channels),
-        nn.ReLU(),
-        nn.MaxPool2d(kernel_size=2, stride=2)
-    )
-
-
-def functional_conv_block(x, weights, biases, bn_weights, bn_biases):
+    if activation == 'relu':
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+    elif activation == 'selu':
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),
+            nn.SELU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+    else:
+        raise ValueError('Unsupported activation.')
+
+
+def functional_conv_block(x, weights, biases, bn_weights, bn_biases, activation: str = 'relu'):
     """Performs 3x3 convolution, ReLu activation, 2x2 max pooling in a functional fashion."""
     x = F.conv2d(x, weights, biases, padding=1)
-    x = F.batch_norm(x, running_mean=None, running_var=None, weight=bn_weights, bias=bn_biases, training=True)
-    x = F.relu(x)
+    if activation == 'relu':
+        x = F.batch_norm(x, running_mean=None, running_var=None, weight=bn_weights, bias=bn_biases, training=True)
+        x = F.relu(x)
+    elif activation == 'selu':
+        x = F.selu(x)
+    else:
+        raise ValueError('Unsupported activation.')
     x = F.max_pool2d(x, kernel_size=2, stride=2)
     return x
 
 
 ##########
 # Models #
 ##########
-def get_few_shot_encoder(num_input_channels=1):
+def get_few_shot_encoder(num_input_channels=1, activation: str = 'relu'):
     """Creates a few shot encoder as used in Matching and Prototypical Networks
 
     # Arguments:
         num_input_channels: Number of color channels the model expects input data to contain. Omniglot = 1,
             miniImageNet = 3
+        activation: Whether to use ReLu activation + batchnorm or SELU on its own
     """
     return nn.Sequential(
-        conv_block(num_input_channels, 64),
-        conv_block(64, 64),
-        conv_block(64, 64),
-        conv_block(64, 64),
+        conv_block(num_input_channels, 64, activation),
+        conv_block(64, 64, activation),
+        conv_block(64, 64, activation),
+        conv_block(64, 64, activation),
         Flatten(),
     )
 
 
 class FewShotClassifier(nn.Module):
-    def __init__(self, num_input_channels: int, k_way: int, final_layer_size: int = 64):
+    def __init__(self, num_input_channels: int, k_way: int, final_layer_size: int = 64, activation: str = 'relu'):
         """Creates a few shot classifier as used in MAML.
 
         This network should be identical to the one created by `get_few_shot_encoder` but with a
         clasification layer on top.
 
-            # Arguments:
-                num_input_channels: Number of color channels the model expects input data to contain. Omniglot = 1,
-                    miniImageNet = 3
-                k_way: Number of classes the model will discriminate between
-                final_layer_size: 64 for Omniglot, 1600 for miniImageNet
+        # Arguments:
+            num_input_channels: Number of color channels the model expects input data to contain. Omniglot = 1,
+                miniImageNet = 3
+            k_way: Number of classes the model will discriminate between
+            final_layer_size: 64 for Omniglot, 1600 for miniImageNet
+            activation: Whether to use ReLu activation + batchnorm or SELU on its own
         """
         super(FewShotClassifier, self).__init__()
-        self.conv1 = conv_block(num_input_channels, 64)
-        self.conv2 = conv_block(64, 64)
-        self.conv3 = conv_block(64, 64)
-        self.conv4 = conv_block(64, 64)
+        self.conv1 = conv_block(num_input_channels, 64, activation)
+        self.conv2 = conv_block(64, 64, activation)
+        self.conv3 = conv_block(64, 64, activation)
+        self.conv4 = conv_block(64, 64, activation)
 
         self.logits = nn.Linear(final_layer_size, k_way)
 
@@ -116,7 +132,7 @@ def functional_forward(self, x, weights):
 
         for block in [1, 2, 3, 4]:
             x = functional_conv_block(x, weights[f'conv{block}.0.weight'], weights[f'conv{block}.0.bias'],
-                                      weights[f'conv{block}.1.weight'], weights[f'conv{block}.1.bias'])
+                                      weights.get(f'conv{block}.1.weight'), weights.get(f'conv{block}.1.bias'))
 
         x = x.view(x.size(0), -1)