How to set batch size when extracting feature using pretained models


(ytian) #1

I am new to pytorch. I want to extract feature using pretrained models, eg. vgg16. I can extract feature from one image at one time. However, I want to improve the speed, I want to set the input batch to 64, rather than 1(which is by default).
How can I set the input batch size? I do not need to fine tune the model, just want to extract feature.
I know this is a useful link, How to extract features of an image from a trained model, but I can not find how to set the batch size here.


(Vahid Mirjalili) #2

You should define a Dataset class, and a then pass that to a Dataloader. This allows you to specify the batch_size.

A dataset class has the following template:

from torch.utils import data

class MyDataset(data.Dataset):
    def __init__(self):
          pass 

    def __getitem__(self, index):
          ## instructions to return a tensor for one image 
          ## and its label (if needed)
          pass

     def __len__(self):
          ## return the size of dataset
          pass
         


(ytian) #3

Thank you, that is helpful.