import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from torchvision.models import resnet18
from torch.backends.cudnn import benchmark
import time
import gc

start_time = None
torch.backends.cudnn.benchmark = True

batch_size = 512
in_size = 4096
out_size = 4096
num_layers = 3
num_batches = 50
epochs = 3
data = [torch.randn(batch_size, in_size, device='cuda') for _ in range(num_batches)]
targets = [torch.randn(batch_size, out_size, device='cuda') for _ in range(num_batches)]

def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.time()


def end_timer_and_print(local_msg):
    torch.cuda.synchronize()
    end_time = time.time()
    print('\n' + local_msg)
    print('Total execution time = {:.3f} sec'.format(end_time -start_time))
    print('Max memory used by tensors = {} bytes'.format(torch.cuda.max_memory_allocated()))


def make_model(in_size, out_size, num_layers):
    layers = []
    for _ in range(num_layers - 1):
        layers.append(nn.Linear(in_size, in_size))
        layers.append(nn.ReLU())
    layers.append(nn.Linear(in_size, out_size))
    return nn.Sequential(*tuple(layers)).cuda()

def train_default():
    net = make_model(in_size, out_size, num_layers)
    opt = torch.optim.SGD(net.parameters(), lr=0.001)
    loss_fn = torch.nn.MSELoss().cuda()
    net.train()
    start_timer()
    for epoch in range(epochs):
        for input, target in zip(data, targets):
            output = net(input)
            loss = loss_fn(output, target)
            loss.backward()
            opt.step()
            opt.zero_grad()
    end_timer_and_print('Default precision:')

def train_mixed():
    net = make_model(in_size, out_size, num_layers)
    opt = torch.optim.SGD(net.parameters(), lr=0.001)
    loss_fn = torch.nn.MSELoss().cuda()
    net.train()
    scaler = GradScaler()
    start_timer()
    for epoch in range(epochs):
        for input, target in zip(data, targets):
            with autocast():
                output = net(input)
                loss = loss_fn(output, input)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
            opt.zero_grad()
    end_timer_and_print('Mixed Precision:')

class SingleConv(nn.Module):
    def __init__(self):
        super(SingleConv, self).__init__()
        self.conv = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)

    def forward(self, x):
        return self.conv(x)

def train_single_conv():
    default_net = SingleConv()
    default_net.cuda()
    default_net.train()
    batches = 5
    default_opt = torch.optim.SGD(default_net.parameters(), lr=0.0001)
    loss_fn = torch.nn.MSELoss().cuda()
    data = [torch.randn(64, 3, 224, 224, device='cuda') for _ in range(batches)]
    tar = [torch.randn(64, 64, 112, 112, device='cuda') for _ in range(batches)]
    start_timer()
    for _ in range(5):
        for input, target in zip(data, tar):
            out = default_net(input)
            loss = loss_fn(out, target)
            loss.backward()
            default_opt.step()
            default_opt.zero_grad()
    end_timer_and_print('Default 1 Conv train:')

    mixed_net = SingleConv()
    mixed_net.cuda()
    mixed_net.train()
    loss_fn = nn.MSELoss().cuda()
    scaler = GradScaler()
    mixed_opt = torch.optim.SGD(mixed_net.parameters(), lr=0.0001)
    start_timer()
    for _ in range(5):
        for input, target in zip(data, tar):
            with autocast():
                out = mixed_net(input)
                loss = loss_fn(out, target)
            scaler.scale(loss).backward()
            scaler.step(mixed_opt)
            scaler.update()
            mixed_opt.zero_grad()
    end_timer_and_print('Mixed 1 Conv train:')

the output is:
Default precision:
Total execution time = 1.760 sec
Max memory used by tensors = 1367458816 bytes
Mixed Precision:
Total execution time = 0.833 sec
Max memory used by tensors = 1401006592 bytes
Default 1 Conv train:
Total execution time = 0.090 sec
Max memory used by tensors = 2470273536 bytes
Mixed 1 Conv train:
Total execution time = 0.127 sec
Max memory used by tensors = 2644665344 bytes