Profile torch.sum (torch.Tensor.sum) with different shape arrangements

torch_sum_prof.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 Pi-Yueh Chuang <pychuang@gwu.edu>
#
# Distributed under terms of the MIT license.

"""
Profile torch.sum (torch.Tensor.sum) with different shape arrangements.
"""
import time
import torch

shapes = [
    [8192, 400, 2],
    [8192, 2, 400],
    [400, 8192, 2],
    [400, 2, 8192],
    [2, 8192, 400],
    [2, 400, 8192]
]

idxs = [2, 1, 2, 1, 0, 0]

profs = []

print('\n'+'='*80+"\nStart profiling\n"+'='*80+'\n')

for i, (shape, idx) in enumerate(zip(shapes, idxs)):

    print("Case {}: Creating tensor with shape {}".format(i, shape))
    A = torch.rand(*shape).to(torch.float64).to("cuda")

    print("Case {}: Performaing warm-ups".format(i))
    for _ in range(10):
        C = torch.sum(A, idx)

    print("Case {}: Performing profiling".format(i))
    with torch.autograd.profiler.profile(use_cuda=True) as prof:
        for _ in range(1000):
            C = torch.sum(A, idx)

    print("Case {}: Output tensor shape - {}".format(i, C.shape))

    profs.append(prof.total_average())
    del A, C
    time.sleep(2)

    print('\n'+'='*80+'\n')

print("Done profiling\n")

s = "Case {}: Matrix shape - {}, CPU time - {}, CUDA time - {}"
for i, (shape, prof) in enumerate(zip(shapes, profs)):
    print(s.format(i, shape, prof.cpu_time_str, prof.cuda_time_str))