# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/secutron/Practice_Ignite/blob/main/A_2_1.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="UfVeNJIVYVu9"
# ## HandlersTimeProfiler

# + [markdown] id="VK6Pe16EpvP3"
# ###A.2.1. HandlersTimeProfiler
#
#  이그나이트의 [HnadlerTimeProfiler](https://pytorch.org/ignite/generated/ignite.contrib.handlers.time_profilers.HandlersTimeProfiler.html)를 이용해 데이터 로딩과 처리, 엔진에 적용된 각종 핸들러에 대한 간단한 프로파일링을 진행하여 아래와 같이 출력하거나 저장할 수 있다. 이그나이트에서 직접 제공하는 간단한 컨볼루션 네트워크에 대한 프로파일링 예제는 [이곳](https://github.com/pytorch/ignite/blob/9230a7319047b37ce19d956e024fa1b86030c30a/examples/notebooks/HandlersTimeProfiler_MNIST.ipynb)에서 참조하도록 한다.
#
#  상기 예제의 실행 결과는 다음과 같다
#
#   <div align="center">
# <img width=800 src="https://i.imgur.com/KObXAVR.png"/>
# </div>
#

# + [markdown] id="HLBKKwFyBI0h"
# 결과 항목 중 제일 아래를 보면 Dataflow에 소요된 시간이 약 101초 임을 알 수 있는데, HandlersTimeProfiler 에서는 미니배치를 로딩하는 이벤트 발생 전과 후에 기록한 시간을 이용하여 이를 측정한다. 이와 마찬가지로 Processing에 소요된 시간은 이터레이션(iteration) 시작과 종료 이벤트 시점에 기록한 시간을 이용하여 측정하며, 예제에서는 약 25초가 소요되었음을 알 수 있다.
#

# + id="HdhoeHGgsimM"
import os

gpu_gtg = False
if int(os.environ.get("COLAB_GPU")) > 0:
    gpu_gtg = "COLAB_GPU" in os.environ

tpu_gtg = "COLAB_TPU_ADDR" in os.environ

if tpu_gtg: # tpu
    print("TPU")
    #VERSION = "nightly"

    # https://github.com/pytorch/builder/pull/750
    VERSION = "20210304" # was 20200607" 

    # !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
    # !python pytorch-xla-env-setup.py --version $VERSION

# + colab={"base_uri": "https://localhost:8080/"} id="27ND-qg2tJEU" outputId="2586824f-1ce4-4390-c5ea-0cd4df8d3df6"
# !pip install --pre pytorch-ignite

# + id="bYoQTy8JtJBU"
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms
from torchvision import datasets, models

import torchsummary

import ignite
import ignite.distributed as idist
from ignite.engine import Engine, Events, create_supervised_evaluator, create_supervised_trainer
from ignite.metrics import Accuracy, Loss, RunningAverage, ConfusionMatrix
from ignite.handlers import ModelCheckpoint, EarlyStopping
from ignite.utils import setup_logger

# + id="y3akwem4zQQf"
from ignite.contrib.handlers import ProgressBar, HandlersTimeProfiler


# + id="va-GVsDLyijz"
def training(local_rank, config, **kwargs):
    print("local rank: ", local_rank)

    ###########################################################
    # 데이터 준비
    train_transform = transforms.Compose(
        [
            transforms.Pad(4),
            transforms.RandomCrop(32, fill=128),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ]
    )

    test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),])

    if idist.get_local_rank() > 0:
        idist.barrier()

    trainset = torchvision.datasets.CIFAR10(root=config["data_path"], train=True, download=True, transform=train_transform)
    testset = torchvision.datasets.CIFAR10(root=config["data_path"], train=False, download=True, transform=test_transform)

    if idist.get_local_rank() == 0:
        idist.barrier()

    trainloader = idist.auto_dataloader(trainset, batch_size=config["batch_size"], shuffle=True, num_workers=config["num_workers"], drop_last=True)
    testloader = idist.auto_dataloader(testset, batch_size=config["batch_size"], shuffle=False, num_workers=config["num_workers"],)

    ###########################################################
    # 모델, 옵티마이저, 로스, 트레이너, 이밸류에이터
    num_classes = 10
    model = models.resnet18(num_classes = num_classes)
       
    model = idist.auto_model(model)
    optimizer = idist.auto_optim(optim.Adam(model.parameters(), lr=0.001))

    criterion = nn.CrossEntropyLoss().to(idist.device())

    trainer = create_supervised_trainer(model, optimizer, criterion, device=idist.device())
    trainer.logger = setup_logger("hkim-trainer")




    #*************************************************************
    # HandlersTimeProfiler
    profiler = HandlersTimeProfiler()
    profiler.attach(trainer)

    # Init and attach progressbar
    pbar = ProgressBar(persist=True)
    pbar.attach(trainer, metric_names="all")




    metrics = {
        'accuracy':Accuracy(),
        'ce':Loss(criterion),
    }

    val_evaluator = create_supervised_evaluator(model, metrics=metrics, device=idist.device())
    val_evaluator.logger = setup_logger("hkim-val_evaluator")

    # track a running average of the scalar loss output for each batch.
    RunningAverage(output_transform=lambda x: x).attach(trainer, 'loss')

    ###########################################################
    # 이벤트

    @trainer.on(Events.EPOCH_COMPLETED)
    def log_validation_results(trainer):
        state = val_evaluator.run(testloader)
        metrics = val_evaluator.state.metrics
        accuracy = metrics['accuracy']*100
        loss = metrics['ce']
        #log_metrics(val_evaluator.logger, state.epoch, state.times["COMPLETED"], "validation evaluator", state.metrics)
        pbar.log_message("Validation Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}".format(state.epoch, accuracy, loss))
        pbar.n = pbar.last_print_n = 0

    trainer.run(trainloader, max_epochs=config["num_epochs"])    


    #*************************************************************
    # HandlersTimeProfiler
    profiler.print_results(profiler.get_results())    

# + colab={"base_uri": "https://localhost:8080/", "height": 1000, "referenced_widgets": ["8aed36fc448e4a4d93d51571dfc27a76", "<KEY>", "126b5fe1aa1f4fb8a6e6efeb8d2e36ff", "5877f880ed104ed79bd1b7054838237b", "<KEY>", "82490e9176c44accaa67bb681ed0dfbf", "866e6260d3ad4e98b32e1f72643f47b7", "633ceb41c520422e857ba4138bd6f468", "<KEY>", "<KEY>", "cc9a9fcadd3f41dbb9f305ed916d71b7", "ffb8b4be834848168f5ea106ec0a09f9", "<KEY>", "<KEY>", "2d10aa4dd6734b76865679a203432259", "c41a97ff37624a69ac1411e1ae7956a4", "de1504db28584a2f8943127b97342ac4", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "ac055ace583645bd89ef2635c653066c", "dcf3ae824de840dfabb011c7b49abed0", "0e7d23d510d54400b32211d14ed2d69c", "c58f79c223ef4ef18fe122da641e2b28", "829d616a99c04347b609ef2743a2727c", "<KEY>", "<KEY>", "7c438a6511d44ca783022fa2c5ce3c48", "e60f72d76bee43c4bfd89804edff644c", "2b436fdb7114401a8009affca4bea300", "<KEY>", "e56e6b31f62e45e48e27d81c9b6dacd8", "ac17c15610424f6cb01ee7433f8567fd", "<KEY>", "5ca4ded0a2934a5993b7bcb467a4d777", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "a49d09561c5a4af4bed759c3a296242e", "482df02980ba425b9e3c4461fb618887", "<KEY>", "<KEY>", "11a65af7dd01461585b8603d59d2db31", "<KEY>", "8be3770eca104a4d81a38cc78d60c39f", "<KEY>", "4e5b66e8fa4d485fa668eec4743087a3"]} id="FXJzm8yqtI-k" outputId="5a893e87-0025-4d13-9a66-e4ebea5ab5aa"
config = {
    "seed": 543,
    "data_path" : "./cifar10",
    "output_path" : "./output-cifar10/",
    "model" : "resnet18",
    "batch_size" : 512,
    "momentum" : 0.9,
    "weight_decay" : 1e-4,
    "num_workers" : 2,
    "num_epochs" : 5,
    "learning_rate" : 0.4,
    "num_warmup_epochs" : 4,
    "validate_every" : 3, 
    "checkpoint_every" : 1000,
    "backend" : None, 
    "resume_from" : None, 
    "log_every_iters" : 15,
    "nproc_per_node" : None, 
    "stop_iteration" : None, 
    "with_amp" : False,
    "log_interval" : 10,
    "verbose_set" : False,
    "verbose_set2" : False,
    "verbose_loader" : False

}

if not (tpu_gtg or gpu_gtg): # cpu
    config["backend"] = 'gloo'
    config["nproc_per_node"] = 8
elif gpu_gtg: # gpu
    config["backend"] = 'nccl'
    config["nproc_per_node"] = 1
elif tpu_gtg: # tpu
    config["backend"] = 'xla-tpu'
    config["nproc_per_node"] = 8
else: # error
    raise RuntimeError("Unknown environment: tpu_gtg {}, gpu_gtg {}".format(tpu_gtg, gpu_gtg))

if config["backend"] == "xla-tpu" and config["with_amp"]:
    raise RuntimeError("The value of with_amp should be False if backend is xla")


dist_configs = {'nproc_per_node': config["nproc_per_node"], "start_method": "fork"}  

#def log_metrics(logger, epoch, elapsed, tag, metrics):
#    metrics_output = "\n".join([f"\t{k}: {v}" for k, v in metrics.items()])
#    logger.info(f"\nEpoch {epoch} - Evaluation time (seconds): {elapsed:.2f} - {tag} metrics:\n {metrics_output}")

with idist.Parallel(backend=config["backend"], **dist_configs) as parallel:
    parallel.run(training, config, a=1, b=1)

# + [markdown] id="3GzqJ_zsDAtS"
#  Processing 시간 항목은, 멀티노드 멀티CPU/GPU/TPU 환경하에 미니배치가 준비된 상태에서 계산되므로 노드 간 또는 GPU간 네트워킹 관련 오버헤드를 배제하고 프로파일 할 수 있다는 점에서 유용하다. 마찬가지로 Dataflow 시간 항목이 크다면 네트워킹 오버헤드가 크다고 생각할 수 있다는 점에서 개발자가 사용하는 시스템에 대한 평가가 일부 가능해진다. 

# + [markdown] id="OQL39wXmXx9k"
# ## License
#

# + [markdown] id="YXio6q3iX5Ig"
# ---
#
#
# Note: This is not an official [LG AI Research](https://www.lgresearch.ai/) product but sample code provided for an educational purpose
#
# <br/>
# author: <NAME>
# <br/>  
# email: <EMAIL> / <EMAIL>  
#
#
# ---

StarcoderdataJupyterScriptsDedupFiltered