Create Pytorch "Stack of Views" to save on GPU memory

Question

I am trying to expand datasets for analysis in Pytorch such that from one 1D (or 2D) tensor two stacks of views are generated. In the following image A (green) and B (blue) are views of the original tensor that are slid from left to right, which would then be combined into single tensors for batch processing:

The motivation behind using views for this is to save on GPU memory, since for large, multi-dimensional datasets this expansion process can convert a dataset of tens of MB into tens of GB despite tremendous data reuse (if normal tensors are used). Simply returning one view at a time is not desirable since the actual processing of tensors works in large batches.

Is what I'm trying to do possible in Pytorch? Simply using torch.stack(list of views) creates a new tensor with a copy of the original data, as verified by tensor.storage().data_ptr().

Another way to phrase the question: can you create batches of tensor views?

The current steps are:

1. Load and pre-process all datasets
2. Convert datasets into tensors and expand into stacks of sliding views, as shown above
3. Move all stacks to GPU to avoid transfer bottleneck during training

Answer 1

As mentioned in the comments, Tensor.unfold can be used for this task. You provide a tensor, starting index, length value, and step size. This returns a batch of views exactly like I was describing, though you have to unfold tensors one at a time for A and B.

The following code can be used to generate A and B:

A = source_tensor[:-B_length].unfold(0, A_length, 1)
B = source_tensor[A_length:].unfold(0, B_length, 1)

A.storage().data_ptr() == source_tensor.storage().data_ptr() returns True

Since the data pointers are the same it is correctly returning views of the original tensor instead of copies, which saves on memory.