'pre-trained network ,CUDA out of memory

    import torch
    from torch import nn
    import torch.nn.functional as F
    import torch.optim as optim
    import matplotlib.pyplot as plt
    %matplotlib inline
    import numpy as np
    import torchvision
    from torchvision import transforms
    import os
    base_dir = './data/4weather'
    train_dir = os.path.join(base_dir,'train')
    test_dir = os.path.join(base_dir,'test')
    Test_transform = transforms.Compose([
        transforms.Resize((192,192)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
    ])
    Train_transform = transforms.Compose([
        transforms.Resize(224),
        transforms.RandomCrop(192),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(0.2),
        #Randomly change the brightness, contrast, saturation and hue of an image.
        transforms.ColorJitter(brightness=0.5),
        transforms.ColorJitter(contrast=0.5),
        transforms.ToTensor(),
        transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
    ])
    train_ds = torchvision.datasets.ImageFolder(train_dir,transform = Train_transform)
    test_ds = torchvision.datasets.ImageFolder(test_dir,transform =Test_transform)
    BATCH_SIZE = 4
    train_dl = torch.utils.data.DataLoader(
        train_ds,
        batch_size = BATCH_SIZE,
        shuffle = True
    )
    test_dl = torch.utils.data.DataLoader(
        test_ds,
        batch_size = BATCH_SIZE
    )
    model = torchvision.models.vgg16(pretrained=True)
    for p in model.features.parameters():
        p.requires_grad = False
    model.classifier[-1].out_features = 4
        model.to('cuda')a
    optimizer = torch.optim.Adam(model.parameters(),lr=0.001)
    loss_fn = nn.CrossEntropyLoss()
    epochs=10
    def train(dl,model,loss_fn,optimizer):
        size = len(dl.dataset)
        num_batches = len(dl)
        train_loss , train_accuracy = 0 , 0
        model.train()
        for x , y in dl:
            x, y  = x.to('cuda') , y.to('cuda')
            pred = model(x)
            loss = loss_fn(pred,y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            with torch.no_grad():
                train_accuracy += (pred.argmax(1)==y).type(torch.float32).sum().item()
                train_loss +=loss.item()
        train_accuracy /=size
        train_loss /= num_batches
        return train_loss,train_accuracy
    def test(dl,model,loss_fn):
        size = len(dl.dataset)
        num_batches = len(dl)
        test_loss , test_accuracy = 0 , 0
        model.eval()
        with torch.no_grad():
            for x , y in dl:
                x, y  = x.to('cuda') , y.to('cuda')
                pred = model(x)
                loss = loss_fn(pred,y)
                test_accuracy += (pred.argmax(1)==y).type(torch.float32).sum().item()
                test_loss +=loss.item()
            test_accuracy /=size
            test_loss /= num_batches
            return test_loss,test_accuracy
    def fit(epochs,train_dl,test_dl,model,loss_fn,optimizer):
        fit_epochs = epochs
        epoch_train_accuracy = []
        epoch_train_loss = []
        epoch_test_accuracy = []
        epoch_test_loss = []
        for epoch in range(fit_epochs):
            train_loss,train_acc = train(train_dl,model,loss_fn,optimizer)
            test_loss,test_acc = test(test_dl,model,loss_fn)
            epoch_train_loss.append(train_loss)
            epoch_train_accuracy.append(train_acc)
            epoch_test_loss.append(test_loss)
            epoch_test_accuracy.append(test_acc)
            template = ('Epoch:{} train_loss: {:.5f},train_accuracy: {:.2f},test_loss: 
                  {:.5f},test_accuracy: {:.2f}.')
            print(template.format(epoch+1,train_loss,train_acc*100,test_loss,test_acc*100))
        print('Done')[`enter image description here`][1]
        return epoch_train_accuracy,epoch_train_loss,epoch_test_accuracy,epoch_test_loss
    (train_accuracy_result,
    train_loss_result,
    test_accuracy_result,
    test_loss_result) = fit(epochs,train_dl,test_dl,model,loss_fn,optimizer)

When I use the pre-trained network VGG16 with PyTorch, I get CUDA out of memory, and no matter how I debug BATCH_SIZE and memory, the result is still the same. This problem does not occur when the pretrained network is not used:

CUDA out of memory. Tried to allocate 392.00 MiB (GPU 0; 4.00 GiB total capacity; 2.28 GiB already allocated; 93.08 MiB free; 2.30 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source