Hi again,
I’m trying to train a pretrained vgg-16 on the Fashion MNIST dataset. I added an extra conv layer at the beggining and modified the last layers to end up to 10 nodes. Also, I applied nn.LogSoftmax(dim=1) at the end, because I’m using nn.NLLLoss(). I disabled gradients for all gradients on model.features, except the first layer, so only the first conv layer and model.classifier layers are trained. I ran my code in two PCs one with 5gb GPU memory and one with 6gb, and both failed. Then, I tried Google Collab on GPU and still it takes too long to run even one epoch. This is my code:
if self.model_type in ['vgg-16', 'vgg-19']:

            

            assert self.image_size == 224, "ERROR: Wrong image size."

            model = torchvision.models.vgg16(pretrained=True) if self.model_type == 'vgg-16' else torchvision.models.vgg19(pretrained=True)

            if self.input_ch != 3:

                first_conv_layer = [nn.Conv2d(self.input_ch, 3, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True)]

                first_conv_layer.extend(list(model.features))  

                model.features= nn.Sequential(*first_conv_layer)  

            model.classifier[-1] = nn.Linear(4096, 1000)

            model.classifier.add_module('7', nn.ReLU())

            model.classifier.add_module('8', nn.Dropout(p=0.5, inplace=False))

            model.classifier.add_module('9', nn.Linear(1000, self.output_ch))

            model.classifier.add_module('10', nn.LogSoftmax(dim=1))

            

            for param in model.features[1:].parameters(): # disable grad for trained layers

                param.requires_grad = False

Is this normal, or am I doing something wrong?
UPDATE: For batch_size=64 it took more than 10 minutes for 1 epoch. Is this normal? I’m going to try a larger batch_size as well to see what happens.
UPDATE 2: For a batch size of 256 I got this error:

RuntimeError: CUDA out of memory. Tried to allocate 784.00 MiB (GPU 0; 15.90 GiB total capacity; 14.87 GiB already allocated; 323.88 MiB free; 14.88 GiB reserved in total by PyTorch)

Wow, 14.88 GB already reserved? Isn’t that too much?