Speed comparison between Torchvision and Torchaug#

Setting for comparison#

We used a python script that compare the speed between Torchvision (v0.15.2) and Torchaug on CPU and GPU.

It does not test the improvement based on eliminating the CPU/GPU synchronization which should favor Torchaug in comparison with Torchvision.

It was performed on 1 GPU A100 and 8 logical cores of an Intel(R) Xeon(R) Silver 4215R CPU @ 3.20GHz.

Batch Transforms#

Input torch.randn(batch_size, 3, 224, 224).

Using Torchvision, to have random call on each element of a batch, one has to iterate over the elements of the batch as the following:

import torchvision.transforms as transforms

t = transforms.ColorJitter(0.5, 0.5, 0.5, 0.1)
torch.stack([t(img) for img in batch_input])

In our implementation, there is no need to iterate over the elements, but some augmentations allow for flexibility between making several random calls for sub-batches and time of computations.

Time in ms averaged on 100 runs for Torchvision and Torchaug.

CPU#

Transform

Rand Calls

Torchvision

Torchaug

Torchvision

Torchaug

Torchvision

Torchaug

Torchvision

Torchaug

Batch size

128

128

64

64

16

16

8

8

BatchRandomColorJitter

1

905.15 ± 44.71

336.76 ± 14.02

62.27 ± 4.38

34.39 ± 3.32

BatchRandomColorJitter

8

506.70 ± 12.96

292.49 ± 5.25

130.44 ± 4.18

67.85 ± 4.03

BatchRandomColorJitter

-1

1475.54 ± 120.92

667.68 ± 7.53

726.49 ± 76.45

374.77 ± 4.12

172.52 ± 41.09

131.69 ± 3.59

85.26 ± 30.58

68.11 ± 3.69

BatchRandomGaussianBlur

753.40 ± 58.82

338.49 ± 5.34

371.74 ± 46.65

145.64 ± 3.91

90.75 ± 22.07

31.79 ± 1.60

46.71 ± 17.19

22.60 ± 2.79

BatchRandomGrayScale

156.61 ± 5.88

47.87 ± 3.81

41.52 ± 4.70

11.06 ± 4.44

7.94 ± 2.06

2.22 ± 3.16

3.97 ± 3.51

2.55 ± 3.27

BatchRandomHorizontalFlip

138.29 ± 4.64

44.62 ± 4.55

25.71 ± 3.57

9.56 ± 3.53

3.34 ± 3.46

1.59 ± 2.83

1.62 ± 2.81

0.64 ± 1.90

BatchRandomResizedCrop

1

233.59 ± 21.54

72.43 ± 12.90

10.36 ± 3.31

4.95 ± 3.43

BatchRandomResizedCrop

8

131.33 ± 4.29

49.83 ± 3.23

16.89 ± 2.28

16.16 ± 1.23

BatchRandomResizedCrop

16

133.11 ± 4.21

48.70 ± 2.86

32.52 ± 2.18

16.27 ± 1.38

BatchRandomResizedCrop

-1

303.01 ± 4.69

299.61 ± 4.04

112.53 ± 2.72

111.21 ± 1.86

26.41 ± 3.25

26.18 ± 3.18

13.24 ± 3.29

13.01 ± 3.37

BatchRandomSolarize

185.82 ± 7.84

103.43 ± 3.95

53.91 ± 5.30

23.78 ± 4.10

10.67 ± 3.80

3.51 ± 3.51

5.41 ± 3.48

1.83 ± 2.96

GPU#

Transform

Rand Calls

Torchvision

Torchaug

Torchvision

Torchaug

Torchvision

Torchaug

Torchvision

Torchaug

Batch size

128

128

64

64

16

16

8

8

BatchRandomColorJitter

1

2.79 ± 0.24

1.52 ± 0.01

1.20 ± 0.01

1.21 ± 0.09

BatchRandomColorJitter

8

7.30 ± 0.03

7.39 ± 0.31

7.39 ± 0.04

3.87 ± 0.03

BatchRandomColorJitter

-1

61.89 ± 5.92

20.47 ± 0.28

31.47 ± 3.63

17.42 ± 0.07

7.79 ± 1.88

7.41 ± 0.04

4.01 ± 1.38

3.88 ± 0.05

BatchRandomGaussianBlur

15.82 ± 1.39

5.43 ± 0.04

7.75 ± 1.01

2.91 ± 0.01

1.95 ± 0.44

0.99 ± 0.01

1.09 ± 0.40

0.71 ± 0.01

BatchRandomGrayScale

5.37 ± 0.36

0.32 ± 0.00

2.74 ± 0.28

0.19 ± 0.00

0.70 ± 0.14

0.17 ± 0.00

0.38 ± 0.10

0.17 ± 0.00

BatchRandomHorizontalFlip

2.15 ± 0.11

0.28 ± 0.00

1.08 ± 0.06

0.18 ± 0.00

0.30 ± 0.03

0.11 ± 0.00

0.18 ± 0.02

0.11 ± 0.00

BatchRandomResizedCrop

1

0.99 ± 0.16

0.55 ± 0.04

0.25 ± 0.01

0.22 ± 0.01

BatchRandomResizedCrop

8

1.13 ± 0.02

1.10 ± 0.01

1.09 ± 0.01

1.11 ± 0.02

BatchRandomResizedCrop

16

2.11 ± 0.14

2.09 ± 0.02

2.13 ± 0.11

1.11 ± 0.02

BatchRandomResizedCrop

-1

12.67 ± 0.35

12.34 ± 0.16

6.31 ± 0.06

6.14 ± 0.07

1.62 ± 0.04

1.58 ± 0.03

0.83 ± 0.01

0.82 ± 0.01

BatchRandomSolarize

3.03 ± 0.23

0.41 ± 0.01

1.56 ± 0.25

0.24 ± 0.00

0.42 ± 0.09

0.15 ± 0.00

0.24 ± 0.06

0.15 ± 0.00

Transforms#

Input torch.randn(3, 224, 224).

Time in ms averaged on 1000 runs for Torchvision and Torchaug.

CPU#

Transform

Torchvision

Torchaug

Normalize

0.48 ± 1.74

0.41 ± 1.54

RandomColorJitter

11.46 ± 11.44

7.99 ± 8.22

RandomGaussianBlur

5.59 ± 6.57

5.67 ± 6.13

RandomSolarize

0.57 ± 1.82

0.59 ± 1.84

GPU#

Transform

Torchvision

Torchaug

Normalize

0.14 ± 0.11

0.07 ± 0.07

RandomColorJitter

1.29 ± 25.39

0.43 ± 0.42

RandomGaussianBlur

0.18 ± 0.18

0.26 ± 0.24

RandomSolarize

0.04 ± 0.02

0.04 ± 0.03