scan_interv.py

"""
Description:
Author: Jiaqi Gu (jqgu@utexas.edu)
Date: 2021-10-24 16:07:39
LastEditors: Jiaqi Gu (jqgu@utexas.edu)
LastEditTime: 2021-10-24 16:07:39
"""

#!/usr/bin/env python
# coding=UTF-8
import argparse
import logging
import os
from typing import Iterable

import mlflow
import numpy as np
import torch
import torch.nn as nn
from pyutils.config import configs
from pyutils.general import ensure_dir
from pyutils.general import logger as lg
from pyutils.plot import set_ms
from pyutils.torch_train import (
    count_parameters,
    get_learning_rate,
    load_model,
    set_torch_deterministic,
)
from pyutils.typing import Criterion, DataLoader, Optimizer, Scheduler

from core import builder

from core.models.devices.mrr_configs import lambda_res, radius_list
from core.models.layers.utils import (
    CrosstalkScheduler,
    GlobalTemperatureScheduler,
    PhaseVariationScheduler,
    calculate_grad_hessian,
)

set_ms()

logging.getLogger("matplotlib.font_manager").disabled = True


def train(
    model: nn.Module,
    train_loader: DataLoader,
    optimizer: Optimizer,
    scheduler: Scheduler,
    epoch: int,
    criterion: Criterion,
    device: torch.device,
) -> None:
    model.train()
    step = epoch * len(train_loader)
    correct = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device, non_blocking=True)
        target = target.to(device, non_blocking=True)

        optimizer.zero_grad()

        output = model(data)

        pred = output.data.max(1)[1]
        correct += pred.eq(target.data).cpu().sum()

        classify_loss = criterion(output, target)

        loss = classify_loss

        loss.backward()

        optimizer.step()
        step += 1

        if batch_idx % configs.run.log_interval == 0:
            lg.info(
                "Train Epoch: {} [{:7d}/{:7d} ({:3.0f}%)] Loss: {:.4f} Class Loss: {:.4f}".format(
                    epoch,
                    batch_idx * len(data),
                    len(train_loader.dataset),
                    100.0 * batch_idx / len(train_loader),
                    loss.data.item(),
                    classify_loss.data.item(),
                )
            )
            mlflow.log_metrics({"train_loss": loss.item()}, step=step)

    scheduler.step()
    accuracy = 100.0 * correct.float() / len(train_loader.dataset)
    lg.info(f"Train Accuracy: {correct}/{len(train_loader.dataset)} ({accuracy:.2f})%")
    mlflow.log_metrics(
        {"train_acc": accuracy.data.item(), "lr": get_learning_rate(optimizer)},
        step=epoch,
    )


def validate(
    model: nn.Module,
    validation_loader: DataLoader,
    epoch: int,
    criterion: Criterion,
    loss_vector: Iterable,
    accuracy_vector: Iterable,
    device: torch.device,
    average_times: int = 1,
) -> None:
    model.eval()
    val_loss, correct = 0, 0
    total_data = 0
    total_batch = 0
    with torch.no_grad():
        for idx in range(average_times):
            for data, target in validation_loader:
                data = data.to(device, non_blocking=True)
                target = target.to(device, non_blocking=True)
                output = model(data)

                val_loss += criterion(output, target).data.item()
                pred = output.data.max(1)[1]
                correct += pred.eq(target.data).cpu().sum()
                total_data += len(data)
                total_batch += 1

    val_loss /= total_batch
    loss_vector.append(val_loss)

    accuracy = 100.0 * correct.float() / total_data
    accuracy_vector.append(accuracy.item())

    lg.info(
        "\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n".format(
            val_loss, correct, total_data, accuracy
        )
    )
    mlflow.log_metrics(
        {"val_acc": accuracy.data.item(), "val_loss": val_loss}, step=epoch
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("config", metavar="FILE", help="config file")
    # parser.add_argument('--run-dir', metavar='DIR', help='run directory')
    # parser.add_argument('--pdb', action='store_true', help='pdb')
    args, opts = parser.parse_known_args()

    configs.load(args.config, recursive=True)
    configs.update(opts)

    if torch.cuda.is_available() and int(configs.run.use_cuda):
        torch.cuda.set_device(configs.run.gpu_id)
        device = torch.device("cuda:" + str(configs.run.gpu_id))
        torch.backends.cudnn.benchmark = True
    else:
        device = torch.device("cpu")
        torch.backends.cudnn.benchmark = False

    if configs.run.deterministic == True:
        set_torch_deterministic()

    model = builder.make_model(device)
    print(model)
    train_loader, validation_loader = builder.make_dataloader()
    criterion = builder.make_criterion().to(device)

    lg.info(f"Number of parameters: {count_parameters(model)}")

    model_name = f"{configs.model.name}_wb-{configs.quantize.weight_bit}_ib-{configs.quantize.input_bit}_icalg-{configs.ic.alg}_icadapt-{configs.ic.adaptive}_icbest-{configs.ic.best_record}"
    checkpoint = f"./checkpoint/{configs.checkpoint.checkpoint_dir}/{model_name}_{configs.checkpoint.model_comment}.pt"

    lg.info(f"Current checkpoint: {checkpoint}")

    mlflow.set_experiment(configs.run.experiment)
    experiment = mlflow.get_experiment_by_name(configs.run.experiment)
    mlflow.start_run(run_name=model_name)
    mlflow.log_params(
        {
            "exp_name": configs.run.experiment,
            "exp_id": experiment.experiment_id,
            "run_id": mlflow.active_run().info.run_id,
            "inbit": configs.quantize.input_bit,
            "wbit": configs.quantize.weight_bit,
            "init_lr": configs.optimizer.lr,
            "ic_alg": configs.ic.alg,
            "ic_adapt": configs.ic.adaptive,
            "ic_best_record": configs.ic.best_record,
            "checkpoint": checkpoint,
            "restore_checkpoint": configs.checkpoint.restore_checkpoint,
            "pid": os.getpid(),
        }
    )
    lg.info(
        f"Experiment {configs.run.experiment} ({experiment.experiment_id}) starts. Run ID: ({mlflow.active_run().info.run_id}). PID: ({os.getpid()}). PPID: ({os.getppid()}). Host: ({os.uname()[1]})"
    )

    try:
        lg.info(configs)
        load_model(
            model,
            configs.checkpoint.restore_checkpoint,
            ignore_size_mismatch=int(configs.checkpoint.no_linear),
        )
        lg.info("Validate pre-trained model (MODE = weight)...")
        validate(model, validation_loader, -3, criterion, [], [], device)
        loss = 0
        for data, target in validation_loader:
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            output = model(data)

            loss = criterion(output, target)
            loss.backward()

        if configs.noise.PV_schedule == "low":
            mean_schedule_fn = lambda x: 0.005 * x
            std_schedule_fn = lambda x: 0.004 * x + 0.002

        elif configs.noise.PV_schedule == "high":
            mean_schedule_fn = lambda x: 0.025 * x
            std_schedule_fn = lambda x: 0.01 * x + 0.005

        phase_variation_scheduler = PhaseVariationScheduler(
            size=[4, 4, 8, 8],
            T_max=20000,
            mean_schedule_fn=mean_schedule_fn,
            std_schedule_fn=std_schedule_fn,
            smoothing_kernel_size=5,
            smoothing_factor=0.05,
            smoothing_mode="arch",
            min_std=0.001,
            momentum=0.9,
            noise_scenario_src=configs.noise.noise_scenario_src,
            noise_scenario_tgt=configs.noise.noise_scenario_tgt,
            random_state=0,
            device=device,
        )

        if configs.noise.TD_schedule == "linear":
            # schedule = "Spatial"
            TD_schedule_fn = lambda x: x + 300
        elif configs.noise.TD_schedule == "cosine":
            # schedule = "perturbation"
            TD_schedule_fn = lambda x: (-0.25 * np.cos(10 * x) + 0.25) + 300
        elif configs.noise.TD_schedule == "uneven":
            TD_schedule_fn = lambda x: x**3 + 300
        else:
            raise NotImplementedError

        global_temperature_scheduler = GlobalTemperatureScheduler(
            size=[4, 4, 8, 8],
            T_max=20000,
            schedule_fn=TD_schedule_fn,
            T0=300,
            lambda_res=lambda_res,
            L_list=2 * np.pi * radius_list,
            hotspot_mode=getattr(configs.noise, "TD_hotspot_mode", "uniform"),
            device=torch.device("cuda"),
        )
        crosstalk_scheduler = CrosstalkScheduler(
            Size=[4, 4, 8, 8],
            crosstalk_coupling_factor=configs.noise.crosstalk_factor,
            interv_h=configs.noise.inter_h,
            interv_v=configs.noise.inter_v,
            device=device,
        )

        model.set_noise_schedulers(
            scheduler_dict={
                "phase_variation_scheduler": phase_variation_scheduler,
                "global_temp_scheduler": global_temperature_scheduler,
                "crosstalk_scheduler": crosstalk_scheduler,
            }
        )

        model.set_phase_variation(configs.noise.set_PV)
        if configs.noise.set_PV:
            lg.info(
                "Validate converted pre-trained model (MODE = phase) with phase variation..."
            )
            validate(model, validation_loader, -1, criterion, [], [], device)

        # set global temp variation
        model.set_global_temp_drift(configs.noise.set_GTD)
        if configs.noise.set_GTD:
            lg.info(
                "Validate converted pre-trained model (MODE = phase) with global temperature variation..."
            )
            validate(model, validation_loader, -1, criterion, [], [], device)

        model.set_crosstalk_noise(configs.noise.set_Crosstalk)

        if configs.noise.set_Crosstalk:
            lg.info(
                "Validate converted pre-trained model (MODE = phase) with crosstalk noise..."
            )
            validate(model, validation_loader, -1, criterion, [], [], device)
        calculate_grad_hessian(
            model, train_loader, criterion, num_samples=10, device=device
        )
        model.backup_ideal_weights()
        model.gen_weight_salience(mode=configs.mapping.salience_mode)
        acc_vectors_noremap = []
        acc_vectors_none = []
        acc_vectors_first = []
        acc_vectors_second = []
        delta_step = configs.mapping.probe_interval
        cycles_vector = []
        total_cycles = 0
        total_cali_cycles = 0
        remap_thres = configs.mapping.W_thres
        T_thres = configs.mapping.T_thres
        counter = 0
        cool_time = configs.mapping.cool_time  # configs.mapping.cool_time
        # probe_num = 0
        probe_cycle = model.probe_cycles()
        inference_cycle = model.cycles(
            [
                1,
                configs.dataset.in_channel,
                configs.dataset.img_height,
                configs.dataset.img_width,
            ]
        )

        remapping_cycles = model.remap(
            input_shape=[
                1,
                configs.dataset.in_channel,
                configs.dataset.img_height,
                configs.dataset.img_width,
            ],
            flag=True,
            alg="LAP",
            salience_mode="first_grad",
            average_times=1,
            tolerance=100,
            verbose=True,
            enable_step=True,
        )
        mapping_cycles = model.map_to_hardware(
            [
                1,
                configs.dataset.in_channel,
                configs.dataset.img_height,
                configs.dataset.img_width,
            ],
            lr=configs.mapping.lr,
            num_steps=configs.mapping.num_steps,
            stop_thres=configs.mapping.stop_thres,
            average_times=configs.mapping.average_times,
            criterion=configs.mapping.criterion,
            # validation_callback=validate_callback_fn,
            sparsity=configs.mapping.sparsity,
            sparsity_mode=configs.mapping.sparsity_mode,
        )
        total_cali_cycles += mapping_cycles + remapping_cycles
        total_cycles += mapping_cycles + remapping_cycles

        for step in range(10000):
            model.set_enable_remap(True)
            model.gen_weight_salience(mode="first_grad")

            if counter == 0 and (step % delta_step == 0 or step == 10000 - 1):
                lg.info(
                    f"step {step}, T: {global_temperature_scheduler.T, global_temperature_scheduler.T0} K"
                )
                # model.set_enable_remap(False)
                # validate(model, validation_loader, -1, criterion, [], acc_vectors_noremap, device, average_times=3)
                remap_flag = bool(
                    (
                        (
                            global_temperature_scheduler.T
                            - global_temperature_scheduler._last_T
                        )
                        * global_temperature_scheduler.hotspot_map
                    ).mean()
                    > T_thres
                )
                print(f"temperature flag is? ={remap_flag}")
                if remap_flag is False:
                    remap_flag = model.probe_weight_error() > remap_thres
                    total_cycles += probe_cycle

                print(f"probed weight_noise is = {model.probe_weight_error()}")
                print(f"remap? ={remap_flag}")
                if remap_flag:
                    remapping_cycles = model.remap(
                        input_shape=[
                            1,
                            configs.dataset.in_channel,
                            configs.dataset.img_height,
                            configs.dataset.img_width,
                        ],
                        flag=True,
                        alg="LAP",
                        salience_mode="first_grad",
                        average_times=1,
                        tolerance=100,
                        verbose=True,
                        enable_step=True,
                    )
                    mapping_cycles = model.map_to_hardware(
                        [
                            1,
                            configs.dataset.in_channel,
                            configs.dataset.img_height,
                            configs.dataset.img_width,
                        ],
                        lr=configs.mapping.lr,
                        num_steps=configs.mapping.num_steps,
                        stop_thres=configs.mapping.stop_thres,
                        average_times=configs.mapping.average_times,
                        criterion=configs.mapping.criterion,
                        # validation_callback=validate_callback_fn,
                        sparsity=configs.mapping.sparsity,
                        sparsity_mode=configs.mapping.sparsity_mode,
                    )
                    total_cali_cycles += mapping_cycles + remapping_cycles
                    total_cycles += mapping_cycles + remapping_cycles

                    counter = cool_time

            counter = max(0, counter - 1)

            # remap_flag = False
            if step % 50 == 0 or step == 10000 - 1:
                validate(
                    model,
                    validation_loader,
                    -1,
                    criterion,
                    [],
                    acc_vectors_first,
                    device,
                    average_times=3,
                )
                lg.info(f"step now = {step}")
                lg.info(
                    f"temp now is {global_temperature_scheduler.T}, previous temp is {global_temperature_scheduler._last_T}"
                )
                lg.info(
                    f"total cali + remap cycles now is {total_cali_cycles}({total_cali_cycles/total_cycles:.4f}), total cycles={total_cycles}"
                )

            model.step_noise_scheduler(T=1)
            total_cycles += inference_cycle
            # total_cycles += cycles
            cycles_vector.append(total_cycles)
        # print(f'number of probeing = {probe_num}')
        # print(acc_vectors_noremap)
        # print(acc_vectors_none)
        lg.info(acc_vectors_first)
        # print(acc_vectors_second)

        ensure_dir(
            f"./Experiment/log/{configs.dataset.name}/{configs.model.name}/remap_interv"
        )
        filename = f"./Experiment/log/{configs.dataset.name}/{configs.model.name}/remap_mode/remap_acc.csv"
        np.savetxt(
            filename,
            np.array(
                [
                    cycles_vector,
                    acc_vectors_noremap,
                    acc_vectors_none,
                    acc_vectors_first,
                    acc_vectors_second,
                ]
            ).T,
            delimiter=",",
        )

    except KeyboardInterrupt:
        lg.warning("Ctrl-C Stopped")


if __name__ == "__main__":
    main()