GPU vectorization via PyTorch

cpu_vs_gpu_torch.py

import numpy as np
import torch
from datetime import datetime


if __name__ == '__main__':
    num_dims = 3
    dim = 1024

    for _ in range(10):
        tensor1 = np.random.randn(*(num_dims * (dim,)))
        tensor2 = np.random.randn(*(num_dims * (dim,)))

        start = datetime.now()
        tensor3 = np.multiply(tensor1, tensor2)
        print(f"Multiplication on CPU took: {datetime.now() - start} seconds")

        tensor1 = torch.Tensor(tensor1)
        tensor2 = torch.Tensor(tensor2)
        tensor1.cuda()
        tensor2.cuda()

        start = datetime.now()
        tensor4 = torch.multiply(tensor1, tensor2)
        print(f"Multiplication on GPU took: {datetime.now() - start} seconds")
        print()

Multiplication on CPU took: 0:00:01.447265 seconds
Multiplication on GPU took: 0:00:00.280588 seconds

Multiplication on CPU took: 0:00:01.447959 seconds
Multiplication on GPU took: 0:00:00.342412 seconds

Multiplication on CPU took: 0:00:01.453325 seconds
Multiplication on GPU took: 0:00:00.349863 seconds

Multiplication on CPU took: 0:00:01.489557 seconds
Multiplication on GPU took: 0:00:00.340789 seconds

Multiplication on CPU took: 0:00:01.439105 seconds
Multiplication on GPU took: 0:00:00.383328 seconds

Multiplication on CPU took: 0:00:01.439311 seconds
Multiplication on GPU took: 0:00:00.426975 seconds

Multiplication on CPU took: 0:00:01.440257 seconds
Multiplication on GPU took: 0:00:00.476254 seconds

Multiplication on CPU took: 0:00:01.458992 seconds
Multiplication on GPU took: 0:00:00.363097 seconds

Multiplication on CPU took: 0:00:01.477163 seconds
Multiplication on GPU took: 0:00:00.358958 seconds

Multiplication on CPU took: 0:00:01.477979 seconds
Multiplication on GPU took: 0:00:00.331797 seconds