Multi-GPU Tensor Initialization Question

[xsys-technology]

The documentation advises the usage of 'type_as' when initializing new tensors in multi-gpu settings:

https://pytorch-lightning.readthedocs.io/en/stable/advanced/multi_gpu.html#init-tensors-using-type-as-and-register-buffer
When you need to create a new tensor, use type_as. This will make your code scale to any arbitrary number of GPUs or TPUs with Lightning.

The example shows a case where a new tensor is initialized inside a LightningModule forward function:

def forward(self, x):
z = torch.Tensor(2, 3)
z = z.type_as(x)

Presumably x is a tensor that has already been initialized on the target gpu.

My question is what to do in the case where we want to initialize a new tensor on the target gpu, and we do not have access to a tensor that has already been initialized on the target gpu?

For example, how does one properly initialize a new tensor when it is created inside a Dataset constructor that is instantiated during LightningDataModule setup()?

class SomeDataModule(LightningDataModule):
...
def setup(self, stage: Optional[str] = None):
if stage in (None, "fit"):
dataset = SomeDataset()
...

where:

class SomeDataset(Dataset):
def init(self):
self.some_tensor = torch.Tensor(2,3)

Will using type_as on the new tensor initialize the data on the target gpu?

self.some_tensor = self.some_tensor.type_as(self.some_tensor)

Or is a different approach necessary? (e.g. register_buffer())

[rohitgr7]

if it's part of the dataset, it's already moved to the target device when a batch is created while iterating over the dataset.

[rohitgr7]

batches are created inside DataLoader's collate_fn. Dataset returns the data based on the individual indices provider by DataLoader's Sampler, which is sent to the collate_fn, where these samples are combined to create a batch.

[xsys-technology]

If things don't seem to be working (it's still going to be a couple hours before I'm ready to test) I'll give self.device a try. This would however seem to go against the advice preparing-your-code.

I was hoping to let lighting automatically handle num_workers correctly, but I'll play around with that if need be.

I was also hoping that lighting will auto-magically handle collate_fn correctly. Again, if need be I can write a custom collate_fn (thank you for reminding me :)

[rohitgr7]

This would however seem to go against the advice preparing-your-code.

Also, those are just suggestions for best practices, especially .cuda, but since here you don't have a reference to a tensor which already exists on the device, it's ok to use .to(device)

I was also hoping that lighting will auto-magically handle collate_fn correctly.

no, you don't have to write a custom collate_fn. PT will handle that as long as these are returned from Dataset get_item.

I was hoping to let lighting automatically handle num_workers correctly,

this is up to the user to set.

[xsys-technology]

Thank you for the good info @rohitgr7 . Seems like I should be able to get things working as planned by using it. I'll post back after I have a chance to test.

[xsys-technology]

After taking a close look at the default collator code, and realizing it did exactly what I wanted, I ended up simply letting it create batches for me by feeding it individual dictionary-of-tensor examples returned through get_item() . All tensors are initialized efficiently (I didn't have to add my own .to(device)) and my gpu's are crunching away with a consistent 96% utilization with a dozen worker threads feeding them. In short, I'm happy :)