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.
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
Thank you, that is helpful.