train_basic.py

# -*- coding: utf-8 -*-
"""
created on 4.19
version: 1.0
train.py
"""

# Imports
import os
script_path = os.path.abspath(__file__)
script_dir = os.path.dirname(script_path)
os.chdir(script_dir)
# os.environ['WANDB_API_KEY'] = 'e7a84490fccf4d551013cad7ca58549bb09594f7'
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
from datetime import datetime
from tqdm import tqdm
# import timm
import numpy as np
# Imports from PyTorch.
import torch
from torch import nn, device, no_grad, save
from torch import max as torch_max
import torch.nn.functional as F
from torch.optim import lr_scheduler
# Imports from networks.
from model import resnet, vgg, lenet, mlp, mobilenetv2,vit
# Imports from utils.
from data.dataset import load_dataset
# Imports from networks.
from recovery.noise_aware.noise_inject import InjectForward, InjectWeight, InjectWeightNoise
from recovery.qat.fake_quantize import fake_quantize_prepare
from utils.utils import test_evaluation, train_step, create_optimizer
from utils.earlystopping import EarlyStopping
from InferHardware.sram.convert_sram import convert_to_sram_prepare

# from call_inference import infer_memtorch, infer_aihwkit, infer_MNSIM
# from call_inference import infer_aihwkit, infer_MNSIM
import argparse
import yaml
import wandb
# wandb.login()

def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace

parser = argparse.ArgumentParser()
parser.add_argument('--project_name', default='AnalogAI', type=str, help='name')
parser.add_argument('--architecture', default='vit', type=str, help='name')
parser.add_argument('--dataset', default='cifar10', type=str, help='training dataset')
parser.add_argument('--batch-size', default=512, type=int, help='mini-batch size')
parser.add_argument('-sram_analog_recover',default=False, type=bool, help='whether to recover on an analog platform')
parser.add_argument('-analog_infer',default=False, type=bool, help='whether to infer on an analog platform')

parser.add_argument('--config', type=str, default='exp1.yml', help='Path to the config file')
args = parser.parse_args()

config_dir = './exp/'
with open(config_dir + args.config, 'r') as f:
    config = yaml.load(f, Loader=yaml.FullLoader)
config = dict2namespace(config)

# Device to use
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Path to store datasets
data_dir = os.path.join(os.getcwd(), "data", config.data.dataset)

# Path to store results
if config.recovery.sram.use:
    save_dir = './save_model/' + config.data.architecture + '_sram'
else:
    save_dir = './save_model/' + config.data.architecture + '_base'
model_path = config.data.architecture + '.pth'
save_path = os.path.join(save_dir, model_path)

# # Training parameters
# random_seed = 2024
# np.random.seed(random_seed)
# torch.manual_seed(random_seed)
# if torch.cuda.device_count() > 1:
#     torch.cuda.manual_seed_all(random_seed)
# else:
#     torch.cuda.manual_seed(random_seed)
# # initialize the early_stopping object
    
# early_stopping = EarlyStopping(patience=20, verbose=True)

random_seed = 2024

import random
import numpy as np
import torch

# 设置 NumPy / Python / PyTorch 随机种子
np.random.seed(random_seed)
random.seed(random_seed)
torch.manual_seed(random_seed)
if torch.cuda.device_count() > 1:
    torch.cuda.manual_seed_all(random_seed)
else:
    torch.cuda.manual_seed(random_seed)

# 设置 cuDNN 为确定性行为，防止非确定性卷积等
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# 强制使用确定性算法（若调用了非确定性操作会报错提示）
torch.use_deterministic_algorithms(True)

# initialize the early_stopping object
early_stopping = EarlyStopping(patience=20, verbose=True)


def training_loop(model, 
                  criterion, 
                  optimizer, 
                  train_data, 
                  validation_data, 
                  config,
                  pla_lr_scheduler,
                  device, 
                  print_every=1):
    """Training loop.

    Args:
        model (nn.Module): Trained model to be evaluated
        criterion (nn.CrossEntropyLoss): criterion to compute loss
        optimizer (Optimizer): analog model optimizer
        train_data (DataLoader): Validation set to perform the evaluation
        validation_data (DataLoader): Validation set to perform the evaluation
        epochs (int): global parameter to define epochs number
        print_every (int): defines how many times to print training progress

    Returns:
        nn.Module, Optimizer, Tuple: model, optimizer, and a tuple of
            lists of train losses, validation losses, and test error
    """
    train_losses = []
    valid_losses = []
    test_error = []

    if config.recovery.qat.use:
        model = fake_quantize_prepare(model=model, 
                                      device=device, 
                                      a_bits=config.recovery.qat.a_bits, 
                                      w_bits=config.recovery.qat.w_bits, )

    # Train model
    for epoch in range(0, config.training.epochs):
        # Train_step
        if  config.recovery.noise.act_inject.use:
            print('====>inject forward noise<====')
            noise_a = InjectForward(config.recovery.noise.act_inject.type, 
                                    config.recovery.noise.act_inject.mean, 
                                    config.recovery.noise.act_inject.sigma, 
                                    config.recovery.noise.act_inject.mask)
        else:
            noise_a = None
            
        if  config.recovery.noise.weight_inject.use:
            print('====>inject weight noise<====')
            noise_w = InjectWeightNoise(model, 
                                       noise_level=config.recovery.noise.weight_inject.level)
        else:
            noise_w = None

        model, optimizer, train_loss = train_step(train_data, 
                                                  model,
                                                  model, 
                                                  criterion, 
                                                  optimizer, 
                                                  device,
                                                  config,
                                                  noise_a,
                                                  noise_w)
        train_losses.append(train_loss)

        if epoch % print_every == (print_every - 1):
            # Validate_step
            with torch.no_grad():
                model, valid_loss, error, accuracy = test_evaluation(
                    validation_data, model, criterion, device
                )
                valid_losses.append(valid_loss)
                test_error.append(error)

            print(
                f"{datetime.now().time().replace(microsecond=0)} --- "
                f"Epoch: {epoch}\t"
                f"Train loss: {train_loss:.4f}\t"
                f"Valid loss: {valid_loss:.4f}\t"
                f"Test error: {error:.2f}%\t"
                f"Test accuracy: {accuracy:.2f}%\t"
            )  
        pla_lr_scheduler.step(valid_loss)   
        best_accuracy = early_stopping(accuracy, 
                                       model.state_dict(), 
                                       config.data.architecture,
                                       0,
                                       epoch, 
                                       save_dir)
        if early_stopping.early_stop:
            print("Early stopping")
            break
        # wandb.log({'epoch':epoch, 'accuracy':best_accuracy, 'train_loss':train_loss, 'valid_loss':valid_loss})
    
    # wandb.finish()
    
    return model, optimizer

def select_model(config, in_channels):
    if config.data.architecture == 'resnet18':
        if config.recovery.qat.use:
            from model.resnetQ import ResNet18Q
            model = ResNet18Q(in_channels)
        else:
            # import pdb;pdb.set_trace()
            model = resnet.resnet18(in_channels)
    elif config.data.architecture == 'mobilenet':
        if config.recovery.qat.use:
            from model.mobilenetv2Q import MobileNetV2Q
            model = MobileNetV2Q(num_classes=10)
        else:
            model = mobilenetv2.MobileNetV2(num_classes=10)
    elif config.data.architecture == 'vit':
        if config.recovery.qat.use:
            from model.vitQ import vitQ
            model = vitQ(in_c=in_channels, num_classes= 10, img_size=32, patch=16, dropout=0.1, 
                         num_layers=7, hidden=384, head=12, mlp_hidden=384, is_cls_token=False)
        else:
            model = vit.ViT(in_c=in_channels, num_classes= 10, img_size=32, patch=16, dropout=0.0, 
                            num_layers=7, hidden=384, head=12, mlp_hidden=384, is_cls_token=False)
    return model


def main():
    """Train a PyTorch CNN analog model with dataset (eg. CIFAR10)."""
    # Make sure the directory where to save the results exist.
    # Results include: Loss vs Epoch graph, Accuracy vs Epoch graph and vector data.
    os.makedirs(save_dir, exist_ok=True)

    # Load datasets.
    dataset = load_dataset(data_dir, 
                           config.training.batch_size, 
                           config.data.dataset)
    train_data, validation_data = dataset.load_images(config)

    #----Load the pytorch model------
    in_channels = 3 if config.data.dataset=='cifar10' else 1

    model = select_model(config, in_channels)

    # if torch.cuda.device_count() > 1:
    #     model = nn.DataParallel(model)     
    model.to(device)
    
    if config.recovery.sram.use and config.recovery.sram.load_fp_checkpoint:
        ckpt_path = config.recovery.sram.fp_checkpoint_path
        print(f"[INFO] Loading pretrained float model from {ckpt_path}")
        checkpoint = torch.load(ckpt_path, map_location=device)
        print(f"[DEBUG] checkpoint keys: {checkpoint.keys()}")
        model.load_state_dict(checkpoint, strict=True)
        # print("[DEBUG] conv1.weight norm =", model.conv1.weight.data.norm().item())
        # ckpt = torch.load(ckpt_path)
        # print("float checkpoint conv1 norm =", ckpt['conv1.weight'].norm())
    

    if config.recovery.sram.use:
        model = convert_to_sram_prepare(model=model, 
                                        device=device,
                                        backend='SRAM', 
                                        parallelism=int(config.recovery.sram.parallelism),
                                        error=config.recovery.sram.intensity,)

    optimizer = create_optimizer(model, 
                                 config.training.lr, 
                                 config.training.momentum, 
                                 config.training.weight_decay, 
                                 config.training.optimizer, 
                                 config.recovery.optimizer.sam, 
                                 config.recovery.optimizer.adaptive)

    pla_lr_scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
                                                      factor=0.5,
                                                      patience=10# verbose=True
                                                     )
        
    criterion = nn.CrossEntropyLoss()

    print(f"\n{datetime.now().time().replace(microsecond=0)} --- " f"Started Training")

    # wandb.init(project="AnalogAI", config=config)
    
    # if config.recovery.sram.use:
    #     model = convert_to_sram_prepare(model=model, 
    #                                     device=device,
    #                                     backend='SRAM', 
    #                                     parallelism=int(config.recovery.sram.parallelism),
    #                                     error=config.recovery.sram.error_rate,)

    model, optimizer = training_loop(model, 
                                     criterion, 
                                     optimizer, 
                                     train_data, 
                                     validation_data, 
                                     config,
                                     pla_lr_scheduler,
                                     device,

    )

    print(f"{datetime.now().time().replace(microsecond=0)} --- " f"Completed Network Training")


if __name__ == "__main__":
    # Execute only if run as the entry point into the program
    main()