working commit 2

models/ae.py

@@ -67,11 +67,12 @@ class AE(nn.Module):
 
 
 class SubSpecCAE(nn.Module):
-    def __init__(self, F=20, T=80, norm='batch', activation='relu', dropout_prob=0.25):
+    def __init__(self, F=20, T=80, norm='batch', activation='relu', dropout_prob=0.25, band=0):
         super(SubSpecCAE, self).__init__()
         self.T = T
         self.F = F
         self.activation = activation
+        self.band = band
         Norm = nn.BatchNorm2d if norm == 'batch' else nn.InstanceNorm2d
         Activation = nn.ReLU if activation == 'relu' else nn.LeakyReLU
         self.encoder = nn.Sequential(
@@ -101,7 +102,7 @@ class SubSpecCAE(nn.Module):
         )
 
     def forward(self, x):
-        x = x[:,3,:,].unsqueeze(1)  # select a single supspec
+        x = x[:,self.band,:,].unsqueeze(1)  # select a single supspec
         encoded = self.encoder(x)
         decoded = self.decoder(encoded)
         return decoded, x

models/layers.py

@@ -13,8 +13,9 @@ class Subspectrogram(object):
             sample = sample.reshape(1, *sample.shape)
         # sample shape: 1 x num_mels x num_frames
         sub_specs = []
-        for i in range(0, sample.shape[1], self.hop_size):
-            sub_spec = sample[:, i:i+self.hop_size:,]
+        neg = 0 if self.height == self.hop_size else self.hop_size
+        for i in range(0, sample.shape[1]-neg, self.hop_size):
+            sub_spec = sample[:, i:i+self.height:,]
             sub_specs.append(sub_spec)
         return np.concatenate(sub_specs)