Hi,
I’m trying to train a simple audio classification model on Colab, but my GPU memory (running on a 16GB instance) use keeps expanding and getting out of control every few epochs. Here is the model definition and a minimal snippet of my training code:
class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(out_channels),
            nn.ELU(),
        )

    def forward(self, x):
        x = self.layers(x)
        return x


class AudioClassifier(nn.Module):
    def __init__(self, stereo=True, dropout=0.1):
        super().__init__()
        in_channels = 2 if stereo else 1
        self.spec = MelspectrogramStretch(hop_length=None, 
                                          num_mels=128, 
                                          fft_length=2048, 
                                          norm='whiten', 
                                          stretch_param=[0.4, 0.4])

        self.features = nn.Sequential(*[
            ConvBlock(in_channels=2, out_channels=32, kernel_size=3, stride=1),
            nn.MaxPool2d(3,3),
            nn.Dropout(p=dropout),
            ConvBlock(in_channels=32, out_channels=64, kernel_size=3, stride=1),
            nn.MaxPool2d(4,4),
            nn.Dropout(p=dropout),
            ConvBlock(in_channels=64, out_channels=64, kernel_size=3, stride=1),
            nn.MaxPool2d(4,4),
            nn.Dropout(p=dropout),
        ])
        self.min_len = 80896
        self.gru_hidden_size = 64
        self.gru_layers = 2

        self.rnn = nn.GRU(128, self.gru_hidden_size, num_layers=self.gru_layers) 
        self.ret = nn.Sequential(*[nn.Linear(self.gru_hidden_size,1), nn.Sigmoid()])
  
    
    def modify_lengths(self, lengths):
        def safe_param(elem):
            return elem if isinstance(elem, int) else elem[0]

        for name, layer in self.features.named_children():
            if isinstance(layer, (nn.Conv2d, nn.MaxPool2d)):
                p, k, s = map(safe_param, [layer.padding, layer.kernel_size,layer.stride]) 
                lengths = ((lengths + 2*p - k)//s + 1).long()

        return torch.where(lengths > 0, lengths, torch.tensor(1, device=lengths.device))

    def _many_to_one(self, t, lengths):
        return t[torch.arange(t.size(0)), lengths - 1]

    def init_hidden(self, batch_size, device):
        return torch.zeros(self.gru_layers, batch_size, self.gru_hidden_size, device=device)

    def forward(self, wave, lengths):
        x = wave
        raw_lengths = lengths
        xt = x.float().transpose(1,2)
        xt, lengths = self.spec(xt, raw_lengths)
        xt = self.features(xt)
        lengths = self.modify_lengths(lengths)
        x = xt.transpose(1, -1)

        batch, time = x.size()[:2]
        x = x.reshape(batch, time, -1)
        lengths = lengths.clamp(max=x.shape[1])

        # Handle variable input size
        x_pack = torch.nn.utils.rnn.pack_padded_sequence(x, lengths.clamp(max=x.shape[1]), batch_first=True)
        x_pack, self.hidden = self.rnn(x_pack)
        x, _ = torch.nn.utils.rnn.pad_packed_sequence(x_pack, batch_first=True)
        x = self._many_to_one(x, lengths)
        x = self.ret(x)
        return x



def train():
    for epoch in range(1,epochs+1):
        model.train()
        batch_losses=[]
        
        for batch_idx, batch in enumerate(pbar):
            optimizer.zero_grad()
            wave, lengths, lbl = batch
            model.hidden = model.init_hidden(BATCH_SIZE, device)
            
            pred = model(wave.to(device), lengths.to(device)).squeeze()
            loss = loss_fn(pred, lbl.to(device))
            
            loss.backward()
            batch_losses.append(loss.detach().item())
            optimizer.step()
            
            del loss, pred, wave, lengths, lbl

I have tried deleting the prediction and input variables every loop, made sure I was detaching every variable I keep, and added an init_hidden() function to refresh the hidden states. These have helped me get to around 3-4 epochs before crashing, but it still happens. I’m running a batch size of just 8, and the input audio files are at most 15-20 seconds long.
Is there anything I can do without reducing the batch size even further? Sorry if there are any stupid mistakes there but I’m a CV guy just getting into audio processing. The code from the classifier was heavily inspired by https://github.com/ksanjeevan/crnn-audio-classification 1.