Pytorch1.5_Cuda10.1_timing.py

gistfile1.txt

import random
import os
import numpy as np
import time
import torch
import torch.nn as nn
from apex import amp

def set_seed(seed: int):
    """Set all seeds to make results reproducible (deterministic mode).
       When seed is None, disables deterministic mode.
    :param seed: an integer to your choosing
    """
    if seed is not None:
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.benchmark = True
        np.random.seed(seed)
        random.seed(seed)
        os.environ['PYTHONHASHSEED'] = str(seed)
set_seed(5)




layer1 = nn.Conv2d(3,64, 3, padding=1, bias=False).cuda()
layer2 = nn.Conv2d(3,64, 3, padding=1, bias=False).cuda()
opt2 = torch.optim.Adam(layer2.parameters(), lr=1e-3)
layer2, opt2 = amp.initialize(layer2, opt2, opt_level="O2")
foo = torch.randn((4,3,16,16)).cuda()

_ = layer1(foo)
loss1 = _.mean().backward()

nruns = 10000

fp32_times = [0]
for idx in range(nruns):
    torch.cuda.synchronize()
    start = time.time()
    _ = layer1(foo)
    loss1 = _.mean().backward()
    torch.cuda.synchronize()
    fp32_times.append(time.time() - start)

apex_times = []
_ = layer2(foo)
loss2 = _.mean()
with amp.scale_loss(loss2, opt2) as scaled_loss:
    scaled_loss.backward()

for idx in range(nruns):
    torch.cuda.synchronize()
    start = time.time()
    _ = layer2(foo)
    loss2 = _.mean()
    with amp.scale_loss(loss2, opt2) as scaled_loss:
        scaled_loss.backward()
    torch.cuda.synchronize()
    apex_times.append(time.time() - start)

print("All times in us")
print(f"Pure fp32 time {1e6*sum(fp32_times)/len(fp32_times)}")
print(f"Apex time {1e6*sum(apex_times)/len(apex_times)}")