Slow Training Time with CNN

johnabascal: <p>(Because of forum rules, I cannot post 5 images so I am attaching a Google Drive Link with photos named Figure_x.png. Please refer to these. <a href="https://drive.google.com/open?id=1iqrjsqHd6wZyliSSTdEfAAkonXYI3sky" rel="nofollow noopener">https://drive.google.com/open?id=1iqrjsqHd6wZyliSSTdEfAAkonXYI3sky</a>)</p>
<p>I was encountering some pretty strange runtime issues when training a CNN. For reference, I am training it on a T4 in a Google Cloud Platform VM.</p>
<p>My original architecture was (I used a ResNet style architecture with these layers):</p>
<p><strong>FIGURE 1</strong></p>
<p>I had been training/testing this with different hyperparameters and whatnot since it only took about 30 minutes to an hour to train. All of a sudden, it began taking extremely long so I reduced the architecture to this:</p>
<p><strong>FIGURE 2</strong></p>
<p>I decided to run a profiler on the training loop with only 16 images and these are the results:</p>
<p><strong>FIGURE 3</strong></p>
<p>Since <code>loss.item()</code> and <code>test_loss.item()</code> are only called to gather training and validation statistics, I tried profiling the code without these functions and this happened:</p>
<p><strong>FIGURE 4</strong></p>
<p>Overall, the training loop took the same time to run! This was pretty strange to my boss and me, so thought about it a bit more and came to the conclusion that this might be happening because until <code>loss.item()</code> is called, there are no functions executed in the graph that PyTorch creates.</p>
<p>Next, I decided to look at the CPU, GPU, and GPU memory usage (using nvtop) to see what was going on during this very long <code>inputs.to(device)</code> function call. I did this by incrementing <code>batch_size</code> by 1 from 1 to 16 (at <code>batch_size</code> 17 I got a memory error). CPU and GPU usage remained fairly similar throughout each of the runs, but memory usage had some very strange behavior.</p>
<p><strong>FIGURE 5</strong></p>
<p>As I increased batch size, memory usage increased normally, then drops off around <code>batch_size</code> = 7, then increases normally again until we reach a memory error.</p>
<p>As of now, I am kind of stumped and don’t know what the underlying cause of my slow training and strange memory usage patterns could be. Any help would be greatly appreciated.</p>
<p>Thanks!</p>
<p>This is the training loop:</p>
<pre><code class="lang-auto">def train_model(model, criterion, optimizer, scheduler, num_epochs=100):
    epoch_loss = []
    epoch_acc = []
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
            
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        
        for phase in ['train', 'test']:
            
            if phase == 'train':
                model.train()
                running_loss = 0.0
                
            elif phase == 'test':
                model.eval()   # Set model to evaluate mode
                running_test_loss = 0.0

            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                optimizer.zero_grad()

                if phase == 'train':
                    with torch.set_grad_enabled(phase=='train'):
                        outputs = model(inputs)
                        loss = criterion(labels, outputs)
                        # backward + optimize only if in training phase
                        loss.backward()
                        optimizer.step()
                        # statistics
                        running_loss += loss.item() * inputs.size(0)
                        
                if phase == 'test':
                    
                    with torch.no_grad():
                        outputs = model(inputs)
                        test_loss = criterion(labels, outputs)
                        running_test_loss += test_loss.item() * inputs.size(0)
                

            if phase == 'train':
                scheduler.step()

        epoch_loss.append(running_loss / dataset_sizes['train'])
        epoch_acc.append(running_test_loss/ dataset_sizes['test'])

        print('Train Loss: {:.4f}\nTest Loss {:.4f}'.format(epoch_loss[epoch], epoch_acc[epoch]))
            
        if epoch % 10==9:
            torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': epoch_loss,
            'acc': epoch_acc
            }, os.getcwd() + '/save_model/{}__.pt'.format(epoch))
            
            
    return model, epoch_loss, epoch_acc
</code></pre>