Parallelizing transforms

Hi all,

I have a conceptual question (which I can back up with code if necessary). Let’s say I’m dealing with vector data where part of each vector is a cartesian coordinate (x,y), and I want to write a transform that maps each cartesian coordinate in each vector to its corresponding polar coordinate.

As far as I understand, the common practice is for transforms to operate on samples one-by-one.
Can I write this transform in a vectorized way, i.e. in a way that it transforms a batch of vectors at a time?

Thank you!

You can use broadcasting.

x = torch.randn(10, 2)

r = torch.sqrt(torch.sum(x**2, dim=1))
t = torch.atan2(x[:,1], x[:0])

Dear Kushajveer,

Thank you for your response.

I know how to use broadcasting; that isn’t my question.

My question is, what is the preferred way to write a transform (e.g. the transforms in torchvision that are passed to dataset constructors) such that it operates on a batch of inputs at a time.

It seems that the convention is to apply the transform within the __getitem__() function of the dataset class, but doing it this way means it can’t be parallel.

Oh, I see what you mean. torchvision transforms are just nn.Modules. If you look at the implementations of these transforms like Resize. So you only need to define the forward pass. In your case, the input to the transform is the batch itself and you can implement your transform just like any other nn.Module.