We used a python script that compare the speed between Torchvision and Torchaug on GPU.
It does not test the improvement based on eliminating the CPU/GPU synchronization which should favor Torchaug in comparison with Torchvision.
If you have a better idea on how to compare the two, feel free to open an issue or PR !
Using Torchvision, to have random call on each of the sample in the batch the following code snippet is used:
torch.stack([torchvision_transform(img) for img in batch])We did not use a single call as shown below as this would call the same transform on the whole batch.
torchvision_transform(batch)It was launched on a NVIDIA RTX 2000 Ada GPU with the following input torch.randint(0, 256, (batch_size, 3, 224, 224), dtype=uint8).
| Transform | Num chunks | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | TV (stack) | Torchaug | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Batch size | 256 | 256 | 128 | 128 | 64 | 64 | 32 | 32 | 16 | 16 | 8 | 8 | 4 | 4 | 2 | 2 | 1 | 1 | ||||||||||
| RandomColorJitter | 1 | 28.82 ± 4.02 | 12.20 ± 0.10 | 3.96 ± 0.10 | 1.60 ± 0.06 | 0.81 ± 0.05 | 0.63 ± 0.02 | 0.63 ± 0.01 | 0.60 ± 0.02 | 0.29 ± 0.27 | ||||||||||||||||||
| RandomColorJitter | 8 | 13.56 ± 0.18 | 6.49 ± 0.10 | 4.23 ± 0.06 | 4.16 ± 0.06 | 4.14 ± 0.05 | 2.18 ± 0.03 | 1.16 ± 0.02 | 0.60 ± 0.07 | 0.28 ± 0.27 | ||||||||||||||||||
| RandomColorJitter | 16 | 12.98 ± 0.30 | 8.23 ± 0.10 | 8.12 ± 0.21 | 8.07 ± 0.07 | 4.14 ± 0.05 | 2.17 ± 0.03 | 1.16 ± 0.02 | 0.59 ± 0.02 | 0.29 ± 0.27 | ||||||||||||||||||
| RandomColorJitter | -1 | 59.80 ± 3.52 | 13.66 ± 0.14 | 29.87 ± 2.47 | 11.07 ± 0.09 | 14.87 ± 1.84 | 8.09 ± 0.10 | 7.52 ± 1.25 | 8.05 ± 0.08 | 3.75 ± 0.89 | 4.14 ± 0.06 | 1.93 ± 0.62 | 2.17 ± 0.04 | 1.00 ± 0.45 | 1.16 ± 0.03 | 0.50 ± 0.34 | 0.59 ± 0.01 | 0.27 ± 0.25 | 0.30 ± 0.27 | |||||||||
| RandomGaussianBlur | 36.86 ± 2.52 | 35.60 ± 0.15 | 18.46 ± 1.57 | 17.79 ± 0.07 | 9.22 ± 1.16 | 8.70 ± 0.05 | 4.62 ± 0.75 | 4.16 ± 0.05 | 2.32 ± 0.54 | 2.12 ± 0.02 | 1.19 ± 0.39 | 1.11 ± 0.01 | 0.63 ± 0.27 | 0.63 ± 0.01 | 0.33 ± 0.20 | 0.38 ± 0.01 | 0.17 ± 0.15 | 0.20 ± 0.18 | ||||||||||
| RandomGrayscale | 5.58 ± 0.21 | 1.45 ± 0.05 | 2.78 ± 0.15 | 0.64 ± 0.02 | 1.39 ± 0.10 | 0.32 ± 0.02 | 0.71 ± 0.08 | 0.17 ± 0.01 | 0.37 ± 0.06 | 0.11 ± 0.01 | 0.19 ± 0.04 | 0.10 ± 0.01 | 0.11 ± 0.03 | 0.09 ± 0.01 | 0.06 ± 0.03 | 0.07 ± 0.00 | 0.04 ± 0.02 | 0.04 ± 0.02 | ||||||||||
| RandomHorizontalFlip | 3.22 ± 0.07 | 1.19 ± 0.06 | 1.56 ± 0.05 | 0.50 ± 0.03 | 0.77 ± 0.04 | 0.25 ± 0.01 | 0.39 ± 0.03 | 0.13 ± 0.01 | 0.20 ± 0.02 | 0.09 ± 0.00 | 0.11 ± 0.02 | 0.07 ± 0.00 | 0.06 ± 0.02 | 0.06 ± 0.00 | 0.04 ± 0.01 | 0.05 ± 0.01 | 0.03 ± 0.01 | 0.02 ± 0.01 | ||||||||||
| RandomResizedCrop | 1 | 4.53 ± 0.30 | 2.28 ± 0.16 | 1.07 ± 0.09 | 0.38 ± 0.04 | 0.20 ± 0.02 | 0.13 ± 0.02 | 0.10 ± 0.01 | 0.09 ± 0.02 | 0.09 ± 0.01 | ||||||||||||||||||
| RandomResizedCrop | 8 | 4.37 ± 0.17 | 2.13 ± 0.05 | 1.04 ± 0.03 | 0.63 ± 0.02 | 0.52 ± 0.02 | 0.51 ± 0.02 | 0.29 ± 0.01 | 0.17 ± 0.01 | 0.09 ± 0.01 | ||||||||||||||||||
| RandomResizedCrop | 16 | 4.24 ± 0.12 | 2.17 ± 0.05 | 1.22 ± 0.03 | 0.96 ± 0.02 | 0.94 ± 0.02 | 0.51 ± 0.02 | 0.29 ± 0.02 | 0.17 ± 0.02 | 0.09 ± 0.01 | ||||||||||||||||||
| RandomResizedCrop | -1 | 16.24 ± 1.71 | 14.15 ± 1.12 | 8.02 ± 0.06 | 7.12 ± 0.05 | 4.03 ± 0.06 | 3.59 ± 0.05 | 2.02 ± 0.04 | 1.81 ± 0.03 | 1.02 ± 0.03 | 0.94 ± 0.03 | 0.53 ± 0.02 | 0.50 ± 0.02 | 0.28 ± 0.02 | 0.29 ± 0.01 | 0.15 ± 0.02 | 0.17 ± 0.01 | 0.09 ± 0.02 | 0.08 ± 0.01 | |||||||||
| RandomSolarize | 4.52 ± 0.19 | 1.70 ± 0.06 | 2.24 ± 0.12 | 0.65 ± 0.02 | 1.11 ± 0.08 | 0.30 ± 0.02 | 0.57 ± 0.13 | 0.16 ± 0.01 | 0.29 ± 0.04 | 0.10 ± 0.01 | 0.16 ± 0.03 | 0.08 ± 0.01 | 0.08 ± 0.02 | 0.08 ± 0.01 | 0.05 ± 0.02 | 0.06 ± 0.01 | 0.03 ± 0.01 | 0.03 ± 0.02 |