Optimizing Pytorch implementation of mix-up augmentation - pytorch

Question

So I have this code here for implementing mix-up augmentation. It's incredibly slow and I'm not sure how to make it faster. It seems like there are some operations that are unavoidable and just by nature slow like scaling images by the weight which is 0.5 then summing up each cell seems like a very slow and unavoidable operation. I'm applying this to Reinforcement Learning so I could be augmenting 64 million images, which is why I need it to be a lot faster.
Note: Here's the original author's implementation but I would assume it's equally as slow as it's essentially the same.
import torch
import utils
import os
import torch.nn.functional as F
import torchvision.transforms as TF
import torchvision.datasets as datasets
dataloader = None
data_iter = None
def _load_data(
sub_path: str, batch_size: int = 256, image_size: int = 84, num_workers: int = 16
):
global data_iter, dataloader
for data_dir in utils.load_config("datasets"):
if os.path.exists(data_dir):
fp = os.path.join(data_dir, sub_path)
if not os.path.exists(fp):
print(f"Warning: path {fp} does not exist, falling back to {data_dir}")
dataloader = torch.utils.data.DataLoader(
datasets.ImageFolder(
fp,
TF.Compose(
[
TF.RandomResizedCrop(image_size),
TF.RandomHorizontalFlip(),
TF.ToTensor(),
]
),
),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
)
data_iter = iter(dataloader)
break
if data_iter is None:
raise FileNotFoundError(
"failed to find image data at any of the specified paths"
)
print("Loaded dataset from", data_dir)
def _load_places(batch_size=256, image_size=84, num_workers=16, use_val=False):
partition = "val" if use_val else "train"
sub_path = os.path.join("places365_standard", partition)
print(f"Loading {partition} partition of places365_standard...")
_load_data(
sub_path=sub_path,
batch_size=batch_size,
image_size=image_size,
num_workers=num_workers,
)
def _load_coco(batch_size=256, image_size=84, num_workers=16, use_val=False):
sub_path = "COCO"
print(f"Loading COCO 2017 Val...")
_load_data(
sub_path=sub_path,
batch_size=batch_size,
image_size=image_size,
num_workers=num_workers,
)
def _get_data_batch(batch_size):
global data_iter
try:
imgs, _ = next(data_iter)
if imgs.size(0) < batch_size:
data_iter = iter(dataloader)
imgs, _ = next(data_iter)
except StopIteration:
data_iter = iter(dataloader)
imgs, _ = next(data_iter)
return imgs.cuda()
def load_dataloader(batch_size, image_size, dataset="coco"):
if dataset == "places365_standard":
if dataloader is None:
_load_places(batch_size=batch_size, image_size=image_size)
elif dataset == "coco":
if dataloader is None:
_load_coco(batch_size=batch_size, image_size=image_size)
else:
raise NotImplementedError(
f'overlay has not been implemented for dataset "{dataset}"'
)
def random_mixup(x, dataset="coco"):
"""Randomly overlay an image from Places or COCO"""
global data_iter
alpha = 0.5
load_dataloader(batch_size=x.size(0), image_size=x.size(-1), dataset=dataset)
imgs = _get_data_batch(batch_size=x.size(0)).repeat(1, x.size(1) // 3, 1, 1)
return ((1 - alpha) * (x / 255.0) + (alpha) * imgs) * 255.0

Answer 1

To optimize you need to use the GPU. You need to use PyTorch tensors and operations.
Example of how to do this with PyTorch:
import torch
import torch.nn.functional as F
import torchvision.transforms as TF
import torchvision.datasets as datasets
# Load the data
dataloader = torch.utils.data.DataLoader(
datasets.ImageFolder(
'path/to/data',
TF.Compose(
[
TF.RandomResizedCrop(84),
TF.RandomHorizontalFlip(),
TF.ToTensor(),
]
),
),
batch_size=256,
shuffle=True,
num_workers=16,
pin_memory=True,
)
# Get a batch of data
imgs, _ = next(iter(dataloader))
# Create a tensor of random weights
alpha = torch.rand(imgs.size(0), 1, 1, 1)
# Create a tensor of random indices
indices = torch.randint(0, imgs.size(0), (imgs.size(0),))
# Create a tensor of random images
imgs2 = imgs[indices]
# Mix the images
imgs = (1 - alpha) * imgs + alpha * imgs2
# You can also do this with a single line of code:
imgs = (1 - alpha) * imgs + alpha * imgs[torch.randint(0, imgs.size(0), (imgs.size(0),))]