Pytorch Tip: Yielding Image Sizes

Need to get image sizes using torchvision datasets?

If you are still reading this on medium I moved my blog to

If you try to load images with torchvision you may get an error like this:

May throw and give you an error like this at the bottom of the stack trace:

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 256 and 519 in dimension 2 at /opt/conda/conda-bld/pytorch_1518243271935/work/torch/lib/TH/generic/THTensorMath.c:2897

Basically because each batch is a single tensor (matrix) and pytorch is saying that all the images in a batch need to be the same size. This makes sense. If you had a 2D matrix row 1 could not have 5 columns if row 2 has 7 columns. Every row in a matrix has to have the same number of columns. If you really want to get the sizes using pytorch you can just set a batch_size of 1. That way each image will be its own tensor and you can record/store the sizes. Like so:

I find this useful because pytorch’s image loader seems to load things in a slightly different order than if i just read the files off disk (not to mention its helpful if you are shuffling via a fixed random seed) and this helps me keep my data in order for when I want to check my predictions. It will output something like this:

torch.Size([1, 3, 256, 256])
torch.Size([1, 3, 519, 253])
torch.Size([1, 3, 256, 256])
torch.Size([1, 3, 256, 256])
torch.Size([1, 3, 256, 256])
torch.Size([1, 3, 512, 680])

As you can see images of many sizes using a pytorch dataset!

That’s all for this little tip. If you like machine learning I have a newsletter that covers machine learning and video games called Generation Machine. Be sure to check it out!


Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store