#!/usr/bin/env python3

"""
Example to run pyprof with imagenet models.
"""

import sys
import torch
import torch.nn as nn
import torchvision.models as models
import torch.cuda.profiler as profiler
import argparse

from apex import pyprof
from apex.optimizers import FusedAdam

def parseArgs():
	parser = argparse.ArgumentParser(prog=sys.argv[0], description="Run popular imagenet models.")

	parser.add_argument("-m",
		type=str,
		default="resnet50",
		choices=["alexnet", "densenet121", "densenet161", "densenet169", "densenet201", "googlenet", "mnasnet0_5", "mnasnet0_75", "mnasnet1_0", "mnasnet1_3", "mobilenet_v2", "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnext50_32x4d", "resnext101_32x8d", "wide_resnet50_2", "wide_resnet101_2", "shufflenet_v2_x0_5", "shufflenet_v2_x1_0", "shufflenet_v2_x1_5", "shufflenet_v2_x2_0", "squeezenet1_0", "squeezenet1_1", "vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn", "inception_v3"],
		help="Model.")

	parser.add_argument("-b",
		type=int,
		default=32,
		help="Batch size.")

	parser.add_argument("-o",
		type=str,
		default="adam",
		choices=["adam", "sgd"],
		help="Optimizer.")

	args = parser.parse_args()
	return args

d = {
	"alexnet":				{'H': 224, 'W': 224, 'opts': {}},

	"densenet121":			{'H': 224, 'W': 224, 'opts': {}},
	"densenet161":			{'H': 224, 'W': 224, 'opts': {}},
	"densenet169":			{'H': 224, 'W': 224, 'opts': {}},
	"densenet201":			{'H': 224, 'W': 224, 'opts': {}},

	"googlenet":			{'H': 224, 'W': 224, 'opts': {'aux_logits': False}},

	"mnasnet0_5":			{'H': 224, 'W': 224, 'opts': {}},
	"mnasnet0_75":			{'H': 224, 'W': 224, 'opts': {}},
	"mnasnet1_0":			{'H': 224, 'W': 224, 'opts': {}},
	"mnasnet1_3":			{'H': 224, 'W': 224, 'opts': {}},

	"mobilenet_v2":			{'H': 224, 'W': 224, 'opts': {}},

	"resnet18":				{'H': 224, 'W': 224, 'opts': {}},
	"resnet34":				{'H': 224, 'W': 224, 'opts': {}},
	"resnet50":				{'H': 224, 'W': 224, 'opts': {}},
	"resnet101":			{'H': 224, 'W': 224, 'opts': {}},
	"resnet152":			{'H': 224, 'W': 224, 'opts': {}},

	"resnext50_32x4d":		{'H': 224, 'W': 224, 'opts': {}},
	"resnext101_32x8d":		{'H': 224, 'W': 224, 'opts': {}},

	"wide_resnet50_2":		{'H': 224, 'W': 224, 'opts': {}},
	"wide_resnet101_2":		{'H': 224, 'W': 224, 'opts': {}},

	"shufflenet_v2_x0_5": 	{'H': 224, 'W': 224, 'opts': {}},
	"shufflenet_v2_x1_0": 	{'H': 224, 'W': 224, 'opts': {}},
	"shufflenet_v2_x1_5": 	{'H': 224, 'W': 224, 'opts': {}},
	"shufflenet_v2_x2_0":	{'H': 224, 'W': 224, 'opts': {}},

	"squeezenet1_0":		{'H': 224, 'W': 224, 'opts': {}},
	"squeezenet1_1":		{'H': 224, 'W': 224, 'opts': {}},

	"vgg11":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg11_bn":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg13":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg13_bn":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg16":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg16_bn":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg19":				{'H': 224, 'W': 224, 'opts': {}},
	"vgg19_bn":				{'H': 224, 'W': 224, 'opts': {}},

	"inception_v3":			{'H': 299, 'W': 299, 'opts': {'aux_logits': False}},
	}

def main():
	args = parseArgs()

	pyprof.nvtx.init()
#	pyprof.nvtx.wrap(fused_adam_cuda, 'adam')

	N = args.b
	C = 3
	H = d[args.m]['H']
	W = d[args.m]['W']
	opts = d[args.m]['opts']
	classes = 1000

	net = getattr(models, args.m)
	net = net(**opts).cuda().half()
	net.train()

	x = torch.rand(N, C, H, W).cuda().half()
	target = torch.empty(N, dtype=torch.long).random_(classes).cuda()

	criterion = nn.CrossEntropyLoss().cuda()
	if (args.o == "sgd"):
		optimizer = torch.optim.SGD(net.parameters(), lr = 0.01, momentum=0.9)
	elif (args.o == "adam"):
		optimizer = FusedAdam(net.parameters())
	else:
		assert False

	#Warm up without profiler
	for i in range(2):
		output = net(x)
		loss = criterion(output, target)
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

	with torch.autograd.profiler.emit_nvtx():
		profiler.start()
		output = net(x)
		loss = criterion(output, target)
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()
		profiler.stop()

if __name__ == "__main__":
	main()

torch.rand torch.autograd.profiler.emit_nvtx torch.cuda.profiler.stop torch.cuda.profiler.start torch.empty torch.nn.CrossEntropyLoss