├── bins ├── parse_misclassification.py ├── convert_models.py ├── valid_setup.py ├── draw_image_by_file.py ├── draw_image_by_folder.py ├── vit_model_init.py └── resnet_model_init.py ├── tools ├── loss_function.py ├── mixup.py ├── model_trainer.py └── common_tools.py ├── dataset ├── cifar10.py ├── imagenet.py └── cifar100.py ├── main ├── src_simulation.py ├── src_train.py └── config.py ├── model ├── vgg_asim.py ├── resnet_asim.py └── vision_transformer_asim.py ├── LICENSE ├── examples ├── train │ ├── resnet18_cifar10_qat │ │ └── config.py │ └── resnet18_cifar10_nat │ │ └── config.py └── simulation │ ├── resnet18_imagenet │ └── config.py │ ├── vitb32_cifar10 │ └── config.py │ ├── vitb32_imagenet │ └── config.py │ └── resnet18_cifar10 │ └── config.py ├── module └── basic_module.py └── README.md /bins/parse_misclassification.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parsing misclassified images for root cause investigation. 4 | 5 | """ 6 | import os 7 | import pickle 8 | import shutil 9 | 10 | 11 | def load_pickle(dir): 12 | with open(dir, 'rb') as f: 13 | data = pickle.load(f) 14 | return data 15 | 16 | 17 | def my_mkdir(my_dir): 18 | if not os.path.isdir(my_dir): 19 | os.makedirs(my_dir) 20 | 21 | 22 | if __name__ == '__main__': 23 | path_pkl = r'E:\Machine_Learning\Project\Cifar100_Image_Classification\run\23-08-21_03-02\misc_img_50.pkl' 24 | data_root_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 25 | out_dir = path_pkl[:-4] # Output files directory 26 | error_info = load_pickle(path_pkl) 27 | 28 | for setname, info in error_info.items(): 29 | for imgs_data in info: 30 | label, pred, path_img_rel = imgs_data 31 | path_img = os.path.join(data_root_dir, os.path.basename(path_img_rel)) 32 | img_dir = os.path.join(out_dir, setname, str(label), str(pred)) # Image folder 33 | my_mkdir(img_dir) 34 | shutil.copy(path_img, img_dir) # Copy files 35 | -------------------------------------------------------------------------------- /tools/loss_function.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Customized loss functions. 4 | 5 | """ 6 | import torch 7 | import torch.nn.functional as F 8 | import torch.nn as nn 9 | 10 | 11 | class LabelSmoothLoss(nn.Module): 12 | def __init__(self, smoothing=0.0): 13 | super(LabelSmoothLoss, self).__init__() 14 | self.smoothing = smoothing 15 | 16 | # Step 1: Log Softmax 17 | # Step 2: Distribute weights. Real label = 1 - smoothing, the others = smoothing / (K - 1). 18 | # Step 3: Calculate loss by cross-entropy. 19 | 20 | def forward(self, input, target): 21 | log_p = F.log_softmax(input, dim=-1) 22 | weight = input.new_ones(input.size()) * self.smoothing / (input.size(-1) - 1) 23 | weight.scatter_(-1, target.unsqueeze(-1), (1. - self.smoothing)) 24 | loss = (-weight * log_p).sum(dim=-1).mean() 25 | 26 | return loss 27 | 28 | 29 | if __name__ == '__main__': # Testbench 30 | 31 | output = torch.tensor([[4.0, 5.0, 10.0], [1.0, 5.0, 4.0], [1.0, 15.0, 4.0]]) 32 | label = torch.tensor([2, 1, 1], dtype=torch.int64) 33 | 34 | loss_function = LabelSmoothLoss(0.001) 35 | loss = loss_function(output, label) 36 | 37 | print('CrossEntropy: {}.'.format(loss)) -------------------------------------------------------------------------------- /bins/convert_models.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Convert PyTorch official model to our .pkl checkpoint for pretraining. 4 | 5 | """ 6 | 7 | import torch 8 | 9 | 10 | def convert_official_model_to_custom_format(official_model_path, output_model_path): 11 | # Load the official model weights 12 | official_weights = torch.load(official_model_path) 13 | 14 | # Create a new dictionary for the checkpoint in the format of checkpoint_best.pkl 15 | new_checkpoint = { 16 | 'model_state_dict': official_weights, # Use the official weights for model_state_dict 17 | 'optimizer_state_dict': {}, # Empty optimizer state dict, assuming starting fresh 18 | 'epoch': 0, # Start from epoch 0 19 | 'best_acc': 0.0 # Starting with 0.0 accuracy 20 | } 21 | 22 | # Save the new checkpoint 23 | torch.save(new_checkpoint, output_model_path) 24 | 25 | 26 | if __name__ == "__main__": 27 | # Define paths 28 | official_model_path = r'E:\Machine_Learning\Project\ASiM\run\vit_b_16-c867db91.pth' 29 | output_model_path = r'E:\Machine_Learning\Project\ASiM\run\vit_b_16-c867db91.pkl' 30 | 31 | # Convert and save the weights 32 | convert_official_model_to_custom_format(official_model_path, output_model_path) 33 | -------------------------------------------------------------------------------- /bins/valid_setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Settle validation set of ImageNet in the same form of training dataset. 4 | 5 | """ 6 | 7 | import os 8 | import shutil 9 | import xml.dom.minidom 10 | 11 | 12 | def make_dir(dir): 13 | """ 14 | Create directory if target root does not exist. 15 | """ 16 | if not os.path.isdir(dir): 17 | os.makedirs(dir) 18 | 19 | 20 | if __name__ == "__main__": # Testbench 21 | # 0. Config 22 | # Set abs directory of label file dir (xml file), newly created validation set dir, and val image dir. 23 | 24 | label_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Annotations\CLS-LOC\val' 25 | folder_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 26 | val_img_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\val' 27 | 28 | # 1. List up all label file dir. 29 | names_labels = os.listdir(label_dir) 30 | names_labels = [os.path.join(label_dir, n) for n in names_labels] 31 | 32 | # 2. Loop for copy files. 33 | for idx in names_labels: 34 | # Read filename and label name from corresponding xml file. 35 | dom = xml.dom.minidom.parse(idx) 36 | root = dom.documentElement 37 | target_file = dom.getElementsByTagName('filename') 38 | filename = target_file[0].firstChild.data 39 | target_label = dom.getElementsByTagName('name') 40 | label = target_label[0].firstChild.data 41 | # Set validation image abs dir and copy files. 42 | img_dir = os.path.join(val_img_dir, (filename + '.JPEG')) 43 | label_dir = os.path.join(folder_dir, label) 44 | make_dir(label_dir) 45 | shutil.copy(img_dir, label_dir) 46 | print('Copy {} imgs to {}'.format(filename, label)) 47 | -------------------------------------------------------------------------------- /bins/draw_image_by_file.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Draw test images (sampling) for simulation (CIFAR-100). 4 | 5 | """ 6 | 7 | import os 8 | import shutil 9 | import random 10 | 11 | 12 | def make_dir(dir): 13 | """ 14 | Create directory if target root does not exist. 15 | """ 16 | if not os.path.isdir(dir): 17 | os.makedirs(dir) 18 | 19 | 20 | def copy_img(imgs, root_dir, setname): 21 | """ 22 | Copy image files from original directory to target directory. 23 | :param imgs: Original images abs directory. 24 | :param root_dir: Target root directory. 25 | :param setname: Target set directory name. 26 | """ 27 | data_dir = os.path.join(root_dir, setname) 28 | make_dir(data_dir) 29 | for path_img in imgs: 30 | print(path_img) 31 | shutil.copy(path_img, data_dir) 32 | print('{} dataset, copy {} imgs to {}'.format(setname, len(imgs), data_dir)) 33 | 34 | 35 | if __name__ == "__main__": 36 | # 0. Config 37 | # random_seed = 6666 38 | draw_ratio = 0.01 # Ratio to draw from target path. 39 | 40 | # 1. Read list and shuffle. 41 | # Set root directory and data directory first. 42 | # random.seed(random_seed) 43 | root_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100' 44 | data_dir = os.path.join(root_dir, 'test_images') 45 | name_imgs = [p for p in os.listdir(data_dir) if p.endswith('.jpg')] # List up image names. 46 | path_imgs = [os.path.join(data_dir, name) for name in name_imgs] # Create abs directory for all images. 47 | random.shuffle(path_imgs) # Shuffle images. 48 | 49 | # 2. Divide list into 3 by pre-defined ratio. 50 | draw_breakpoint = int(len(path_imgs) * draw_ratio) 51 | test_imgs = path_imgs[:draw_breakpoint] 52 | 53 | # 3. Copy files. 54 | copy_img(test_imgs, root_dir, 'images_under_test_100') 55 | -------------------------------------------------------------------------------- /tools/mixup.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Mixup Implementation. 4 | 5 | """ 6 | 7 | import numpy as np 8 | import torch 9 | 10 | 11 | def mixup_data(data, label, alpha=1.0, device=True): 12 | """ 13 | Returns mixed inputs, pairs of target, and lambda. 14 | """ 15 | # Get lambda from beta distribution, where the parameter alpha == beta. 16 | lam = np.random.beta(alpha, alpha) if alpha > 0 else 1 17 | 18 | # To get labels from mixed images. 19 | batch_size = data.size()[0] 20 | index = torch.randperm(batch_size).to(device) 21 | 22 | # MixUp. 23 | mixed_img = lam * data + (1 - lam) * data[index, :] 24 | label_a, label_b = label, label[index] 25 | 26 | return mixed_img, label_a, label_b, lam 27 | 28 | 29 | def mixup_loss_function(loss_function, pred, label_a, label_b, lamda): 30 | """ 31 | Mixup integrated into loss function. 32 | """ 33 | return lamda * loss_function(pred, label_a) + (1 - lamda) * loss_function(pred, label_b) 34 | 35 | 36 | if __name__ == '__main__': # Testbench 37 | import cv2 38 | import matplotlib.pyplot as plt 39 | path_1 = r'E:\Machine_Learning\Dataset\Processing\102flowers\jpg\image_00001.jpg' 40 | path_2 = r'E:\Machine_Learning\Dataset\Processing\102flowers\jpg\image_08061.jpg' 41 | 42 | img_1 = cv2.imread(path_1) 43 | img_2 = cv2.imread(path_2) 44 | img_1 = cv2.resize(img_1, (224, 224)) 45 | img_2 = cv2.resize(img_2, (224, 224)) 46 | 47 | alpha = 1. 48 | figsize = 15 49 | plt.figure(figsize=(int(figsize), int(figsize))) 50 | for i in range(1, 10): 51 | lamda = np.random.beta(alpha, alpha) 52 | img_mixup = (img_1 * lamda + img_2 * (1 - lamda)).astype(np.uint8) 53 | img_mixup = cv2.cvtColor(img_mixup, cv2.COLOR_BGR2RGB) 54 | plt.subplot(3, 3, i) 55 | plt.title('lambda = {:.2f}'.format(lamda)) 56 | plt.imshow(img_mixup) 57 | plt.show() 58 | -------------------------------------------------------------------------------- /bins/draw_image_by_folder.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Draw test images (sampling) for simulation (CIFAR10/ImageNet). 4 | 5 | """ 6 | 7 | import os 8 | import shutil 9 | import random 10 | 11 | 12 | def make_dir(dir): 13 | """ 14 | Create directory if target root does not exist. 15 | """ 16 | if not os.path.isdir(dir): 17 | os.makedirs(dir) 18 | 19 | 20 | if __name__ == "__main__": # Testbench 21 | # 0. Config 22 | # Set random seed. 23 | # random_seed = 6666 24 | draw_count = 10000 # Number of images to draw from target path. 25 | 26 | # 1. Read list and shuffle. 27 | # Set root directory and data directory first. 28 | # random.seed(random_seed) 29 | root_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC' 30 | data_dir = os.path.join(root_dir, 'valid') 31 | # Create all subdirectories, mirroring the structure of data_dir 32 | for folder_name in os.listdir(data_dir): 33 | target_subdir = os.path.join(os.path.join(root_dir, 'images_under_test_10000'), folder_name) 34 | make_dir(target_subdir) # This will create both empty and non-empty folders 35 | folder_dir = [p for p in os.listdir(data_dir)] 36 | folder_dir = [os.path.join(data_dir, name) for name in folder_dir] 37 | path_imgs = [] 38 | for dir in folder_dir: 39 | names = [p for p in os.listdir(dir) if p.endswith('.JPEG')] # List up image names. 40 | paths = [os.path.join(dir, name) for name in names] # Create abs directory for all images. 41 | path_imgs += paths 42 | random.shuffle(path_imgs) # Shuffle images. 43 | 44 | # 2. Divide list by pre-defined ratio. 45 | draw_breakpoint = int(draw_count) 46 | test_imgs = path_imgs[:draw_breakpoint] 47 | 48 | test_folder = [p.split('\\')[-2] for p in test_imgs] 49 | target_dir = [os.path.join(root_dir, 'images_under_test_10000', n) for n in test_folder] 50 | 51 | # 3. Copy files. 52 | for draw in range(draw_count): 53 | make_dir(target_dir[draw]) 54 | shutil.copy(test_imgs[draw], target_dir[draw]) 55 | print('Copy {} to {}.'.format(test_imgs[draw], target_dir[draw])) 56 | -------------------------------------------------------------------------------- /dataset/cifar10.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Dataset file of image classification task for . 4 | 5 | """ 6 | 7 | import os 8 | from PIL import Image 9 | from torch.utils.data import Dataset 10 | 11 | 12 | class Cifar10Dataset(Dataset): 13 | 14 | def __init__(self, root_dir, transform=None): 15 | """ 16 | Initialize the dataset with directory and preprocessing methods. 17 | """ 18 | self.root_dir = root_dir 19 | self.transform = transform 20 | self.img_info = [] # [(dir, label), ... , ] 21 | self.class_num = 10 # Number of categories 22 | self.names = tuple(range(self.class_num)) 23 | self._get_img_info() 24 | 25 | def __getitem__(self, index): 26 | """ 27 | Retrieve an item from the dataset at the specified index. 28 | """ 29 | img_dir, label = self.img_info[index] 30 | img = Image.open(img_dir).convert('RGB') 31 | 32 | if self.transform: 33 | img = self.transform(img) 34 | 35 | return img, label, img_dir 36 | 37 | def __len__(self): 38 | """ 39 | Get the total number of items in the dataset. 40 | """ 41 | if not self.img_info: 42 | raise Exception(f'\ndata_dir: {self.root_dir} is an empty directory! Please check your directory settings!') 43 | return len(self.img_info) 44 | 45 | def _get_img_info(self): 46 | """ 47 | Read dataset and summarize the directory & label in a list (img_info). 48 | [(dir, label), ... ,] 49 | """ 50 | for label_num, class_name in enumerate(sorted(os.listdir(self.root_dir))): 51 | class_dir = os.path.join(self.root_dir, class_name) 52 | if not os.path.isdir(class_dir): 53 | continue 54 | 55 | for img_name in sorted(os.listdir(class_dir)): 56 | if img_name.endswith('.png'): 57 | img_path = os.path.join(class_dir, img_name) 58 | self.img_info.append((img_path, label_num)) 59 | 60 | 61 | if __name__ == "__main__": # Testbench 62 | 63 | root_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar10\archive\cifar10\train' 64 | # root_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar10\archive\cifar10\test' 65 | test_dataset = Cifar10Dataset(root_dir) 66 | 67 | print(len(test_dataset)) 68 | print(next(iter(test_dataset))) 69 | -------------------------------------------------------------------------------- /dataset/imagenet.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Dataset file of image classification task for . 4 | 5 | """ 6 | 7 | import os 8 | from PIL import Image 9 | from torch.utils.data import Dataset 10 | 11 | 12 | class ImageNetDataset(Dataset): 13 | 14 | def __init__(self, root_dir, transform=None): 15 | """ 16 | Initialize the dataset with directory and preprocessing methods. 17 | """ 18 | self.root_dir = root_dir 19 | self.transform = transform 20 | self.img_info = [] # [(dir, label), ... , ] 21 | self.class_num = 1000 # Number of categories 22 | self.names = tuple(range(self.class_num)) 23 | self._get_img_info() 24 | 25 | def __getitem__(self, index): 26 | """ 27 | Retrieve an item from the dataset at the specified index. 28 | """ 29 | img_dir, label = self.img_info[index] 30 | img = Image.open(img_dir).convert('RGB') 31 | 32 | if self.transform: 33 | img = self.transform(img) 34 | 35 | return img, label, img_dir 36 | 37 | def __len__(self): 38 | """ 39 | Get the total number of items in the dataset. 40 | """ 41 | if not self.img_info: 42 | raise Exception(f'\ndata_dir: {self.root_dir} is an empty directory! Please check your directory settings!') 43 | return len(self.img_info) 44 | 45 | def _get_img_info(self): 46 | """ 47 | Read dataset and summarize the directory & label in a list (img_info). 48 | [(dir, label), ... ,] 49 | """ 50 | for label_num, class_name in enumerate(sorted(os.listdir(self.root_dir))): 51 | class_dir = os.path.join(self.root_dir, class_name) 52 | if not os.path.isdir(class_dir): 53 | continue 54 | 55 | for img_name in sorted(os.listdir(class_dir)): 56 | if img_name.endswith('.JPEG'): 57 | img_path = os.path.join(class_dir, img_name) 58 | self.img_info.append((img_path, label_num)) 59 | 60 | 61 | if __name__ == "__main__": # Testbench 62 | 63 | # root_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 64 | root_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 65 | test_dataset = ImageNetDataset(root_dir) 66 | 67 | print(len(test_dataset)) 68 | print(next(iter(test_dataset))) 69 | -------------------------------------------------------------------------------- /dataset/cifar100.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Dataset file of image classification task for . 4 | 5 | """ 6 | 7 | import os 8 | import pandas as pd 9 | from PIL import Image 10 | from torch.utils.data import Dataset 11 | 12 | 13 | class Cifar100Dataset(Dataset): 14 | 15 | def __init__(self, root_dir, transform=None): 16 | """ 17 | Initialize the dataset with directory and preprocessing methods. 18 | """ 19 | self.root_dir = root_dir 20 | self.transform = transform 21 | self.img_info = [] # [(dir, label), ... , ] 22 | self.class_num = 100 # Number of categories 23 | self.names = tuple(range(self.class_num)) 24 | self._get_img_info() 25 | 26 | def __getitem__(self, index): 27 | """ 28 | Retrieve an item from the dataset at the specified index. 29 | """ 30 | img_dir, label = self.img_info[index] 31 | img = Image.open(img_dir).convert('RGB') 32 | 33 | if self.transform: 34 | img = self.transform(img) 35 | 36 | return img, label, img_dir 37 | 38 | def __len__(self): 39 | """ 40 | Get the total number of items in the dataset. 41 | """ 42 | if not self.img_info: 43 | raise Exception(f'\ndata_dir: {self.root_dir} is an empty directory! Please check your directory settings!') 44 | return len(self.img_info) 45 | 46 | def _get_img_info(self): 47 | """ 48 | Read dataset and summarize the directory & label in a list (img_info). 49 | """ 50 | # List up image file paths. 51 | names_imgs = [os.path.join(self.root_dir, n) for n in os.listdir(self.root_dir) if n.endswith('.jpg')] 52 | 53 | # Read image IDs & labels from .csv files. 54 | train_label_file = os.path.join(self.root_dir, '..', 'train.csv') 55 | valid_label_file = os.path.join(self.root_dir, '..', 'test.csv') 56 | 57 | try: 58 | df_train = pd.read_csv(train_label_file) 59 | df_valid = pd.read_csv(valid_label_file) 60 | except FileNotFoundError as e: 61 | raise Exception(f'Error reading label files: {e}') 62 | 63 | df = pd.concat([df_train, df_valid]) 64 | content = dict(zip(df['image_id'], df['fine_labels'])) 65 | 66 | # Generate image file path and label pairs within the specified directory. 67 | self.img_info = [(p, content[os.path.basename(p)]) for p in names_imgs if os.path.basename(p) in content] 68 | 69 | 70 | if __name__ == "__main__": # Testbench 71 | 72 | root_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 73 | test_dataset = Cifar100Dataset(root_dir) 74 | 75 | print(len(test_dataset)) 76 | print(next(iter(test_dataset))) 77 | -------------------------------------------------------------------------------- /main/src_simulation.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Main simulation code for the framework. 4 | 5 | """ 6 | 7 | import torch 8 | import numpy as np 9 | from torch.utils.data import DataLoader 10 | from main.config import cfg 11 | from model.resnet_asim import resnet18_asim, resnet34_asim, resnet50_asim, resnet101_asim, resnet152_asim 12 | from model.vgg_asim import vgg11_asim, vgg13_asim, vgg16_asim, vgg19_asim, vgg11_bn_asim, vgg13_bn_asim, vgg16_bn_asim, vgg19_bn_asim 13 | from model.vision_transformer_asim import vit_b_16_asim, vit_b_32_asim, vit_l_16_asim, vit_l_32_asim, vit_h_14_asim 14 | from dataset.cifar10 import Cifar10Dataset 15 | from dataset.cifar100 import Cifar100Dataset 16 | from dataset.imagenet import ImageNetDataset 17 | 18 | if __name__ == '__main__': # Main script for simulation 19 | 20 | # Step 1: Dataset 21 | # Construct Dataset instance, and then construct DataLoader 22 | if cfg.task == 'CIFAR-10': 23 | test_data = Cifar10Dataset(root_dir=cfg.test_dir, transform=cfg.transforms_test) 24 | if cfg.task == 'CIFAR-100': 25 | test_data = Cifar100Dataset(root_dir=cfg.test_dir, transform=cfg.transforms_test) 26 | if cfg.task == 'ImageNet': 27 | test_data = ImageNetDataset(root_dir=cfg.test_dir, transform=cfg.transforms_test) 28 | 29 | test_loader = DataLoader(dataset=test_data, batch_size=cfg.test_bs, shuffle=False, num_workers=cfg.test_workers) 30 | 31 | # Step 2: Model Selection 32 | model_dic = { 33 | 'ResNet-18': resnet18_asim(), 34 | 'ResNet-34': resnet34_asim(), 35 | 'ResNet-50': resnet50_asim(), 36 | 'ResNet-101': resnet101_asim(), 37 | 'ResNet-152': resnet152_asim(), 38 | 'VGG11': vgg11_asim(), 39 | 'VGG13': vgg13_asim(), 40 | 'VGG16': vgg16_asim(), 41 | 'VGG19': vgg19_asim(), 42 | 'VGG11-BN': vgg11_bn_asim(), 43 | 'VGG13-BN': vgg13_bn_asim(), 44 | 'VGG16-BN': vgg16_bn_asim(), 45 | 'VGG19-BN': vgg19_bn_asim(), 46 | 'ViT-B-16': vit_b_16_asim(), 47 | 'ViT-B-32': vit_b_32_asim(), 48 | 'ViT-L-16': vit_l_16_asim(), 49 | 'ViT-L-32': vit_l_32_asim(), 50 | 'ViT-H-14': vit_h_14_asim() 51 | } 52 | model = model_dic[cfg.model_name] 53 | pretrained_state_dict = torch.load(cfg.pretrain_model_path, map_location='cpu') 54 | model.load_state_dict(pretrained_state_dict['model_state_dict']) 55 | model.to(cfg.device) 56 | model.eval() 57 | 58 | # Step 3: Inference (Simulation) 59 | class_num = cfg.cls_num 60 | conf_mat = np.zeros((class_num, class_num)) 61 | 62 | for i, data in enumerate(test_loader): 63 | inputs, labels, path_img = data 64 | inputs, labels = inputs.to(cfg.device), labels.to(cfg.device) 65 | 66 | outputs = model(inputs) 67 | 68 | # Generate confusion matrix 69 | _, predicted = torch.max(outputs.data, 1) 70 | for j in range(len(labels)): 71 | cate_i = labels[j].cpu().numpy() 72 | pre_i = predicted[j].cpu().numpy() 73 | conf_mat[cate_i, pre_i] += 1. 74 | # Simulation print 75 | print('Label: {}, Predict: {}.'.format(cate_i, pre_i)) 76 | 77 | acc_avg = conf_mat.trace() / conf_mat.sum() 78 | print('Test Acc: {:.2%}'.format(acc_avg)) 79 | -------------------------------------------------------------------------------- /tools/model_trainer.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Model Trainer. 4 | 5 | """ 6 | 7 | import torch 8 | import numpy as np 9 | from collections import Counter 10 | from tools.mixup import mixup_data, mixup_loss_function 11 | 12 | 13 | class ModelTrainer(): 14 | 15 | @staticmethod 16 | def train(data_loader, model, loss_f, optimizer, scheduler, epoch_idx, device, cfg, logger): 17 | model.train() 18 | 19 | class_num = data_loader.dataset.class_num 20 | conf_mat = np.zeros((class_num, class_num)) 21 | loss_sigma = [] 22 | loss_mean = 0 23 | acc_avg = 0 24 | dir_misc = [] 25 | label_list = [] 26 | 27 | for i, data in enumerate(data_loader): 28 | 29 | # _, labels = data 30 | inputs, labels, img_dir = data 31 | label_list.extend(labels.tolist()) 32 | 33 | # inputs, labels = data 34 | inputs, labels, img_dir = data 35 | inputs, labels = inputs.to(device), labels.to(device) 36 | 37 | # Mixup 38 | if cfg.mixup: 39 | mixed_inputs, label_a, label_b, lam = mixup_data(inputs, labels, cfg.mixup_alpha, device) 40 | inputs = mixed_inputs 41 | 42 | # forward & backward 43 | outputs = model(inputs) 44 | optimizer.zero_grad() 45 | 46 | # Calculate loss 47 | if cfg.mixup: 48 | loss = mixup_loss_function(loss_f, outputs.cpu(), label_a.cpu(), label_b.cpu(), lam) 49 | else: 50 | loss = loss_f(outputs.cpu(), labels.cpu()) 51 | 52 | loss.backward() 53 | optimizer.step() 54 | 55 | # Record loss 56 | loss_sigma.append(loss.item()) 57 | loss_mean = np.mean(loss_sigma) 58 | 59 | # Generate data for confusion matrix 60 | _, predicted = torch.max(outputs.data, 1) 61 | for j in range(len(labels)): 62 | cate_i = labels[j].cpu().numpy() 63 | pre_i = predicted[j].cpu().numpy() 64 | conf_mat[cate_i, pre_i] += 1. 65 | if cate_i != pre_i: 66 | dir_misc.append((cate_i, pre_i, img_dir[j])) # Record information 67 | acc_avg = conf_mat.trace() / conf_mat.sum() # Correct count on diagonal 68 | 69 | # Print training information every 10 interation. 70 | if i % cfg.log_interval == cfg.log_interval - 1: 71 | logger.info('Training: Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}'. 72 | format(epoch_idx + 1, cfg.max_epoch, i + 1, len(data_loader), loss_mean, acc_avg)) 73 | logger.info('Epoch: {} Sampler: {}'.format(epoch_idx + 1, Counter(label_list))) 74 | 75 | return loss_mean, acc_avg, conf_mat, dir_misc 76 | 77 | @staticmethod 78 | def valid(data_loader, model, loss_f, device): 79 | model.eval() 80 | 81 | class_num = data_loader.dataset.class_num 82 | conf_mat = np.zeros((class_num, class_num)) 83 | loss_sigma = [] 84 | dir_misc = [] 85 | 86 | for i, data in enumerate(data_loader): 87 | # inputs, labels = data 88 | inputs, labels, img_dir = data 89 | inputs, labels = inputs.to(device), labels.to(device) 90 | 91 | outputs = model(inputs) 92 | loss = loss_f(outputs.cpu(), labels.cpu()) 93 | 94 | # Generate data for confusion matrix 95 | _, predicted = torch.max(outputs.data, 1) 96 | for j in range(len(labels)): 97 | cate_i = labels[j].cpu().numpy() 98 | pre_i = predicted[j].cpu().numpy() 99 | conf_mat[cate_i, pre_i] += 1. 100 | if cate_i != pre_i: 101 | dir_misc.append((cate_i, pre_i, img_dir)) # Record information 102 | 103 | # Record loss 104 | loss_sigma.append(loss.item()) 105 | 106 | acc_avg = conf_mat.trace() / conf_mat.sum() # Correct count on the diagonal 107 | 108 | return np.mean(loss_sigma), acc_avg, conf_mat, dir_misc 109 | 110 | 111 | if __name__ == '__main__': 112 | pass 113 | -------------------------------------------------------------------------------- /bins/vit_model_init.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Convert PyTorch official model to our ASiM ViT model checkpoint for pretraining. 4 | 5 | """ 6 | 7 | import sys 8 | import torch 9 | import torch.nn as nn 10 | from main.config import cfg 11 | from model.vision_transformer_asim import vit_b_16_asim, vit_b_32_asim, vit_l_16_asim, vit_l_32_asim, vit_h_14_asim 12 | from module.asim_linear import ASiMLinear 13 | 14 | 15 | def transfer_weights(asim_model, official_weights): 16 | asim_state_dict = asim_model.state_dict() 17 | transferred_layers = 0 18 | total_layers = 0 19 | 20 | # Dictionary of weight block name in PyTorch model to be converted into ASiM model 21 | official_to_asim = { 22 | 'class_token': 'cls_token', 23 | 'conv_proj.weight': 'patch_embedding.conv_proj.weight', 24 | 'conv_proj.bias': 'patch_embedding.conv_proj.bias', 25 | 'encoder.pos_embedding': 'encoder.pos_embedding', 26 | 'heads.head.weight': 'heads.weight', 27 | 'heads.head.bias': 'heads.bias', 28 | 'encoder.ln.weight': 'encoder.ln.weight', 29 | 'encoder.ln.bias': 'encoder.ln.bias', 30 | } 31 | 32 | for i in range(12): 33 | official_to_asim.update({ 34 | f'encoder.layers.encoder_layer_{i}.ln_1.weight': f'encoder.layers.{i}.ln_1.weight', 35 | f'encoder.layers.encoder_layer_{i}.ln_1.bias': f'encoder.layers.{i}.ln_1.bias', 36 | f'encoder.layers.encoder_layer_{i}.self_attention.in_proj_weight': f'encoder.layers.{i}.self_attention.qkv.weight', 37 | f'encoder.layers.encoder_layer_{i}.self_attention.in_proj_bias': f'encoder.layers.{i}.self_attention.qkv.bias', 38 | f'encoder.layers.encoder_layer_{i}.self_attention.out_proj.weight': f'encoder.layers.{i}.self_attention.proj.weight', 39 | f'encoder.layers.encoder_layer_{i}.self_attention.out_proj.bias': f'encoder.layers.{i}.self_attention.proj.bias', 40 | f'encoder.layers.encoder_layer_{i}.ln_2.weight': f'encoder.layers.{i}.ln_2.weight', 41 | f'encoder.layers.encoder_layer_{i}.ln_2.bias': f'encoder.layers.{i}.ln_2.bias', 42 | f'encoder.layers.encoder_layer_{i}.mlp.linear_1.weight': f'encoder.layers.{i}.mlp.mlp1.weight', 43 | f'encoder.layers.encoder_layer_{i}.mlp.linear_1.bias': f'encoder.layers.{i}.mlp.mlp1.bias', 44 | f'encoder.layers.encoder_layer_{i}.mlp.linear_2.weight': f'encoder.layers.{i}.mlp.mlp2.weight', 45 | f'encoder.layers.encoder_layer_{i}.mlp.linear_2.bias': f'encoder.layers.{i}.mlp.mlp2.bias', 46 | }) 47 | 48 | for name, param in official_weights.items(): 49 | total_layers += 1 50 | if name in official_to_asim: 51 | asim_name = official_to_asim[name] 52 | if asim_name in asim_state_dict: 53 | try: 54 | if asim_state_dict[asim_name].shape == param.shape: 55 | asim_state_dict[asim_name].copy_(param) 56 | transferred_layers += 1 57 | print(f"Successfully copied weight for layer: {asim_name}") 58 | else: 59 | print(f"Shape mismatch for layer: {asim_name}, skipping...") 60 | except Exception as e: 61 | print(f"Could not copy weight for layer: {asim_name}. Error: {e}") 62 | else: 63 | print(f"Layer {asim_name} not found in ASiM model.") 64 | else: 65 | print(f"Layer {name} not found in official to ASiM mapping.") 66 | 67 | asim_model.load_state_dict(asim_state_dict) 68 | print(f"Transferred {transferred_layers}/{total_layers} layers successfully.") 69 | sys.stdout.flush() 70 | 71 | 72 | def modify_classifier_head(asim_model, num_classes): 73 | # Modify the classifier head to fit the class number 74 | in_features = asim_model.heads.in_features 75 | asim_model.heads = ASiMLinear(in_features, 76 | num_classes, 77 | bias=True, 78 | wbit=cfg.asim_vit_fc_wbit, 79 | xbit=cfg.asim_vit_fc_xbit, 80 | adc_prec=cfg.asim_vit_fc_adc_prec, 81 | nrow=cfg.asim_vit_fc_nrow, 82 | rand_noise_sigma=cfg.asim_vit_fc_rand_noise_sigma, 83 | non_linear_sigma=cfg.asim_vit_fc_non_linear_sigma, 84 | act_enc=cfg.asim_vit_fc_act_enc, 85 | signed_act=True, 86 | layer='fc', 87 | mode=cfg.asim_vit_fc_mode, 88 | trim_noise=cfg.asim_vit_fc_trim_noise, 89 | device=cfg.device) 90 | nn.init.xavier_uniform_(asim_model.heads.weight) 91 | if asim_model.heads.bias is not None: 92 | nn.init.zeros_(asim_model.heads.bias) 93 | print(f"Modified classifier head to {num_classes} classes.") 94 | sys.stdout.flush() 95 | 96 | 97 | def convert_official_model_to_asim_format(asim_model, num_classes, official_model_path, output_model_path): 98 | # Load the official model weights 99 | official_weights = torch.load(official_model_path) 100 | print(f"Loaded official model weights from {official_model_path}") 101 | sys.stdout.flush() 102 | 103 | # Transfer model weights 104 | transfer_weights(asim_model, official_weights) 105 | 106 | modify_classifier_head(asim_model, num_classes) 107 | 108 | # Create pickle checkpoint 109 | new_checkpoint = { 110 | 'model_state_dict': asim_model.state_dict(), # Converted ASiM model 111 | 'optimizer_state_dict': {}, # Keep blank 112 | 'epoch': 0, # Keep blank 113 | 'best_acc': 0.0 # Keep blank 114 | } 115 | 116 | # Save new checkpoint 117 | torch.save(new_checkpoint, output_model_path) 118 | print(f"Saved new model checkpoint to {output_model_path}") 119 | sys.stdout.flush() 120 | 121 | 122 | if __name__ == "__main__": 123 | # Define official model and output model directory 124 | official_model_path = r'E:\Machine_Learning\Model\vit_b_32-d86f8d99.pth' 125 | output_model_path = r'E:\Machine_Learning\Model\vit_b_32_cls_1000.pkl' 126 | 127 | # Give ASiM model 128 | model = vit_b_32_asim() 129 | num_classes = 1000 130 | 131 | # Convert model 132 | convert_official_model_to_asim_format(model, num_classes, official_model_path, output_model_path) 133 | -------------------------------------------------------------------------------- /main/src_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Main training code for the framework. 4 | 5 | """ 6 | 7 | import sys 8 | import pickle 9 | import torch.optim as optim 10 | import torch.nn as nn 11 | from torch.utils.data import DataLoader 12 | from tools.model_trainer import ModelTrainer 13 | from tools.common_tools import * 14 | from main.config import cfg 15 | from model.resnet_asim import resnet18_asim, resnet34_asim, resnet50_asim, resnet101_asim, resnet152_asim 16 | from model.vgg_asim import vgg11_asim, vgg13_asim, vgg16_asim, vgg19_asim, vgg11_bn_asim, vgg13_bn_asim, vgg16_bn_asim, vgg19_bn_asim 17 | from model.vision_transformer_asim import vit_b_16_asim, vit_b_32_asim, vit_l_16_asim, vit_l_32_asim, vit_h_14_asim 18 | from datetime import datetime 19 | from dataset.cifar10 import Cifar10Dataset 20 | from dataset.cifar100 import Cifar100Dataset 21 | from dataset.imagenet import ImageNetDataset 22 | from tools.loss_function import LabelSmoothLoss 23 | 24 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 25 | sys.path.append(os.path.join(BASE_DIR, '..')) 26 | 27 | if __name__ == '__main__': # Main script for model training 28 | 29 | # Step 0: Config 30 | # Generate Logger 31 | res_dir = os.path.join(BASE_DIR, 'run') 32 | logger, log_dir = make_logger(res_dir) 33 | 34 | # Step 1: Dataset 35 | # Construct Dataset instance, and then construct DataLoader 36 | if cfg.task == 'CIFAR-10': 37 | train_data = Cifar10Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 38 | valid_data = Cifar10Dataset(root_dir=cfg.valid_dir, transform=cfg.transforms_valid) 39 | if cfg.task == 'CIFAR-100': 40 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 41 | valid_data = Cifar100Dataset(root_dir=cfg.valid_dir, transform=cfg.transforms_valid) 42 | if cfg.task == 'ImageNet': 43 | train_data = ImageNetDataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 44 | valid_data = ImageNetDataset(root_dir=cfg.valid_dir, transform=cfg.transforms_valid) 45 | 46 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True, num_workers=cfg.train_workers) 47 | valid_loader = DataLoader(dataset=valid_data, batch_size=cfg.valid_bs, num_workers=cfg.train_workers) 48 | 49 | # Step 2: Model Selection 50 | model_dic = { 51 | 'ResNet-18': resnet18_asim(), 52 | 'ResNet-34': resnet34_asim(), 53 | 'ResNet-50': resnet50_asim(), 54 | 'ResNet-101': resnet101_asim(), 55 | 'ResNet-152': resnet152_asim(), 56 | 'VGG11': vgg11_asim(), 57 | 'VGG13': vgg13_asim(), 58 | 'VGG16': vgg16_asim(), 59 | 'VGG19': vgg19_asim(), 60 | 'VGG11-BN': vgg11_bn_asim(), 61 | 'VGG13-BN': vgg13_bn_asim(), 62 | 'VGG16-BN': vgg16_bn_asim(), 63 | 'VGG19-BN': vgg19_bn_asim(), 64 | 'ViT-B-16': vit_b_16_asim(), 65 | 'ViT-B-32': vit_b_32_asim(), 66 | 'ViT-L-16': vit_l_16_asim(), 67 | 'ViT-L-32': vit_l_32_asim(), 68 | 'ViT-H-14': vit_h_14_asim() 69 | } 70 | model = model_dic[cfg.model_name] 71 | # Load pretrained model 72 | if cfg.pretrain_model_path is not None: 73 | pretrained_state_dict = torch.load(cfg.pretrain_model_path, map_location='cpu') 74 | model.load_state_dict(pretrained_state_dict['model_state_dict']) 75 | logger.info('Load pretrained model.') 76 | model.to(cfg.device) # To device (cpu or gpu) 77 | 78 | # Step 3: Loss function & Optimizer 79 | if cfg.label_smooth: 80 | loss_f = LabelSmoothLoss(cfg.label_smooth_eps) 81 | else: 82 | loss_f = nn.CrossEntropyLoss() 83 | if cfg.model_name in ['ViT-B-16', 'ViT-B-32', 'ViT-L-16', 'ViT-L-32', 'ViT-H-14']: 84 | optimizer = optim.Adam(model.parameters(), lr=cfg.lr_init) 85 | else: 86 | optimizer = optim.SGD(model.parameters(), lr=cfg.lr_init, momentum=cfg.momentum, weight_decay=cfg.weight_decay) 87 | scheduler = optim.lr_scheduler.MultiStepLR(optimizer, gamma=cfg.factor, milestones=cfg.milestones) 88 | 89 | # Step 4: Training iteration 90 | # Record model, loss function, optimizer, and cfg in training 91 | logger.info('cfg:\n{}\n loss_f:\n{}\n scheduler:\n{}\n optimizer:\n{}\n model:\n{}\n'.format( 92 | cfg, loss_f, scheduler, optimizer, model)) 93 | 94 | loss_rec = {'Train': [], 'Valid': []} 95 | acc_rec = {'Train': [], 'Valid': []} 96 | best_acc, best_epoch = 0, 0 97 | for epoch in range(cfg.max_epoch): 98 | 99 | loss_train, acc_train, mat_train, path_error_train = ModelTrainer.train( 100 | train_loader, model, loss_f, optimizer, scheduler, epoch, cfg.device, cfg, logger) 101 | 102 | loss_valid, acc_valid, mat_valid, path_error_valid = ModelTrainer.valid( 103 | valid_loader, model, loss_f, cfg.device) 104 | 105 | logger.info('Epoch[{:0>3}/{:0>3}], Train Acc: {:.2%}, Valid Acc: {:.2%}, Train Loss: {:.4f}, Valid Loss: {:.4f}, LR: {}'.format( 106 | epoch + 1, cfg.max_epoch, acc_train, acc_valid, loss_train, loss_valid, optimizer.param_groups[0]['lr'])) 107 | scheduler.step() 108 | 109 | # Record training information 110 | loss_rec['Train'].append(loss_train), loss_rec['Valid'].append(loss_valid) 111 | acc_rec['Train'].append(acc_train), acc_rec['Valid'].append(acc_valid) 112 | # Confusion matrix 113 | show_confmat(mat_train, train_data.names, 'Train', log_dir, epoch=epoch, verbose=epoch == cfg.max_epoch - 1) 114 | show_confmat(mat_valid, valid_data.names, 'Valid', log_dir, epoch=epoch, verbose=epoch == cfg.max_epoch - 1) 115 | # Save loss, acc transfer curves 116 | plt_x = np.arange(1, epoch + 2) 117 | plot_result(plt_x, loss_rec['Train'], plt_x, loss_rec['Valid'], mode='loss', dir=log_dir) 118 | plot_result(plt_x, acc_rec['Train'], plt_x, acc_rec['Valid'], mode='acc', dir=log_dir) 119 | 120 | # Save model 121 | if best_acc < acc_valid or epoch == cfg.max_epoch - 1: 122 | best_epoch = epoch if best_acc < acc_valid else best_epoch 123 | best_acc = acc_valid if best_acc < acc_valid else best_acc 124 | checkpoint = {'model_state_dict': model.state_dict(), 125 | 'optimizer_state_dict': optimizer.state_dict(), 126 | 'epoch': epoch, 127 | 'best_acc': best_acc} 128 | pkl_name = 'checkpoint_{}.pkl'.format(epoch + 1) if epoch == cfg.max_epoch - 1 else 'checkpoint_best.pkl' 129 | path_checkpoint = os.path.join(log_dir, pkl_name) 130 | torch.save(checkpoint, path_checkpoint) 131 | 132 | # Save directory of misclassified figures 133 | misc_img_name = 'misc_img_{}.pkl'.format(epoch + 1) if epoch == cfg.max_epoch - 1 else 'misc_img_best.pkl' 134 | path_misc_img = os.path.join(log_dir, misc_img_name) 135 | misc_info = {} 136 | misc_info['Train'] = path_error_train 137 | misc_info['Valid'] = path_error_valid 138 | pickle.dump(misc_info, open(path_misc_img, 'wb')) 139 | 140 | logger.info('{} done, Best Acc: {} in epoch {}.'.format( 141 | datetime.strftime(datetime.now(), '%m-%d_%H_%M'), best_acc, best_epoch)) 142 | -------------------------------------------------------------------------------- /bins/resnet_model_init.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Convert PyTorch official model to our ASiM ResNet model checkpoint for pretraining. 4 | 5 | """ 6 | 7 | import sys 8 | import torch 9 | import torch.nn as nn 10 | from module.basic_module import PACT 11 | from module.asim_linear import ASiMLinear 12 | from main.config import cfg 13 | from model.resnet_asim import resnet18_asim, resnet34_asim, resnet50_asim, resnet101_asim, resnet152_asim 14 | 15 | 16 | def transfer_weights(asim_model, official_weights): 17 | asim_state_dict = asim_model.state_dict() 18 | transferred_layers = 0 19 | total_layers = 0 20 | 21 | # Dictionary of weight block name in PyTorch model to be converted into ASiM model 22 | official_to_asim = { 23 | 'bn1.weight': 'bn1.weight', 24 | 'bn1.bias': 'bn1.bias', 25 | 'bn1.running_mean': 'bn1.running_mean', 26 | 'bn1.running_var': 'bn1.running_var', 27 | 'fc.weight': 'fc.weight', 28 | 'fc.bias': 'fc.bias', 29 | } 30 | 31 | # Map convolutional layer based on `cfg.large_model` 32 | if cfg.large_model: 33 | official_to_asim['conv1.weight'] = 'conv2.weight' # Use conv2 if large_model is True 34 | else: 35 | official_to_asim['conv1.weight'] = 'conv1.weight' # Use conv1 otherwise 36 | 37 | # Dynamically update dictionary for layer-by-layer mappings 38 | for i in range(4): 39 | for j in range(2): 40 | official_to_asim.update({ 41 | f'layer{i+1}.{j}.conv1.weight': f'layer{i+1}.{j}.conv1.weight', 42 | f'layer{i+1}.{j}.bn1.weight': f'layer{i+1}.{j}.bn1.weight', 43 | f'layer{i+1}.{j}.bn1.bias': f'layer{i+1}.{j}.bn1.bias', 44 | f'layer{i+1}.{j}.bn1.running_mean': f'layer{i+1}.{j}.bn1.running_mean', 45 | f'layer{i+1}.{j}.bn1.running_var': f'layer{i+1}.{j}.bn1.running_var', 46 | f'layer{i+1}.{j}.conv2.weight': f'layer{i+1}.{j}.conv2.weight', 47 | f'layer{i+1}.{j}.bn2.weight': f'layer{i+1}.{j}.bn2.weight', 48 | f'layer{i+1}.{j}.bn2.bias': f'layer{i+1}.{j}.bn2.bias', 49 | f'layer{i+1}.{j}.bn2.running_mean': f'layer{i+1}.{j}.bn2.running_mean', 50 | f'layer{i+1}.{j}.bn2.running_var': f'layer{i+1}.{j}.bn2.running_var', 51 | }) 52 | 53 | # Add downsample layers mapping if present 54 | if f'layer{i + 1}.{j}.downsample.0.weight' in asim_state_dict: 55 | official_to_asim.update({ 56 | f'layer{i + 1}.{j}.downsample.0.weight': f'layer{i + 1}.{j}.downsample.0.weight', 57 | f'layer{i + 1}.{j}.downsample.1.running_mean': f'layer{i + 1}.{j}.downsample.1.running_mean', 58 | f'layer{i + 1}.{j}.downsample.1.running_var': f'layer{i + 1}.{j}.downsample.1.running_var', 59 | f'layer{i + 1}.{j}.downsample.1.weight': f'layer{i + 1}.{j}.downsample.1.weight', 60 | f'layer{i + 1}.{j}.downsample.1.bias': f'layer{i + 1}.{j}.downsample.1.bias', 61 | }) 62 | 63 | for name, param in official_weights.items(): 64 | total_layers += 1 65 | if name in official_to_asim: 66 | asim_name = official_to_asim[name] 67 | if asim_name in asim_state_dict: 68 | try: 69 | if asim_state_dict[asim_name].shape == param.shape: 70 | asim_state_dict[asim_name].copy_(param) 71 | transferred_layers += 1 72 | print(f"Successfully copied weight for layer: {asim_name}") 73 | else: 74 | print(f"Shape mismatch for layer: {asim_name}, skipping...") 75 | except Exception as e: 76 | print(f"Could not copy weight for layer: {asim_name}. Error: {e}") 77 | else: 78 | print(f"Layer {asim_name} not found in ASiM model.") 79 | else: 80 | print(f"Layer {name} not found in official to ASiM mapping.") 81 | 82 | if cfg.PACT: 83 | # Set default values for PACT 84 | for name, module in asim_model.named_modules(): 85 | if isinstance(module, PACT): 86 | if 'alpha' not in module.__dict__: 87 | module.alpha = nn.Parameter(torch.tensor(6.0)) 88 | print(f"Initialized alpha for {name}") 89 | 90 | asim_model.load_state_dict(asim_state_dict) 91 | print(f"Transferred {transferred_layers}/{total_layers} layers successfully.") 92 | sys.stdout.flush() 93 | 94 | 95 | def modify_classifier_head(asim_model, num_classes): 96 | # Modify the classifier head to fit the class number 97 | in_features = asim_model.fc.in_features 98 | asim_model.fc = ASiMLinear(in_features, num_classes, 99 | wbit=cfg.asim_linear_wbit, 100 | xbit=cfg.asim_linear_xbit, 101 | adc_prec=cfg.asim_linear_adc_prec, 102 | nrow=cfg.asim_linear_nrow, 103 | rand_noise_sigma=cfg.asim_conv_rand_noise_sigma, 104 | non_linear_sigma=cfg.asim_conv_non_linear_sigma, 105 | act_enc=cfg.asim_linear_act_enc, 106 | signed_act=cfg.asim_linear_signed_act, 107 | layer=cfg.asim_linear_layer, 108 | mode=cfg.asim_linear_mode, 109 | trim_noise=cfg.asim_linear_trim_noise, 110 | device=cfg.device) 111 | nn.init.xavier_uniform_(asim_model.fc.weight) 112 | if asim_model.fc.bias is not None: 113 | nn.init.zeros_(asim_model.fc.bias) 114 | print(f"Modified classifier head to {num_classes} classes.") 115 | sys.stdout.flush() 116 | 117 | 118 | def convert_official_model_to_asim_format(asim_model, num_classes, official_model_path, output_model_path): 119 | # Load the official model weights 120 | official_weights = torch.load(official_model_path) 121 | print(f"Loaded official model weights from {official_model_path}") 122 | sys.stdout.flush() 123 | 124 | # Transfer model weights 125 | transfer_weights(asim_model, official_weights) 126 | 127 | # Modify classifier head 128 | modify_classifier_head(asim_model, num_classes) 129 | 130 | # Create pickle checkpoint 131 | new_checkpoint = { 132 | 'model_state_dict': asim_model.state_dict(), # Converted ASiM model 133 | 'optimizer_state_dict': {}, # Keep blank 134 | 'epoch': 0, # Keep blank 135 | 'best_acc': 0.0 # Keep blank 136 | } 137 | 138 | # Save new checkpoint 139 | torch.save(new_checkpoint, output_model_path) 140 | print(f"Saved new model checkpoint to {output_model_path}") 141 | sys.stdout.flush() 142 | 143 | 144 | if __name__ == "__main__": 145 | # Define official model and output model directory 146 | official_model_path = r'E:\Machine_Learning\Model\resnet18-f37072fd.pth' # Replace with your file path 147 | output_model_path = r'E:\Machine_Learning\Model\resnet18_cls_1000.pkl' 148 | 149 | # Instantiate ASiM model 150 | model = resnet18_asim() 151 | num_classes = 1000 152 | 153 | # Convert model 154 | convert_official_model_to_asim_format(model, num_classes, official_model_path, output_model_path) 155 | -------------------------------------------------------------------------------- /model/vgg_asim.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | VGG models for the simulation framework. 4 | 5 | """ 6 | 7 | import torch 8 | import torch.nn as nn 9 | from main.config import cfg 10 | from module.basic_module import QuantConv2d, PACT 11 | from module.asim_conv import ASiMConv2d 12 | from module.asim_linear import ASiMLinear 13 | 14 | 15 | def asim_conv(in_planes, out_planes, stride=1): 16 | return ASiMConv2d(in_planes, 17 | out_planes, 18 | kernel_size=3, 19 | stride=stride, 20 | padding=1, 21 | bias=True, 22 | wbit=cfg.asim_conv_wbit, 23 | xbit=cfg.asim_conv_xbit, 24 | adc_prec=cfg.asim_conv_adc_prec, 25 | nrow=cfg.asim_conv_nrow, 26 | rand_noise_sigma=cfg.asim_conv_rand_noise_sigma, 27 | non_linear_sigma=cfg.asim_conv_non_linear_sigma, 28 | act_enc=cfg.asim_act_enc, 29 | signed_act=cfg.asim_signed_act, 30 | hybrid_levels=cfg.asim_cnn_hybrid_levels, 31 | mode=cfg.asim_conv_mode, 32 | trim_noise=cfg.asim_conv_trim_noise, 33 | device=cfg.device) 34 | 35 | 36 | def asim_affine(in_features, out_features): 37 | return ASiMLinear(in_features, 38 | out_features, 39 | bias=True, 40 | wbit=cfg.asim_linear_wbit, 41 | xbit=cfg.asim_linear_xbit, 42 | adc_prec=cfg.asim_linear_adc_prec, 43 | nrow=cfg.asim_linear_nrow, 44 | rand_noise_sigma=cfg.asim_linear_rand_noise_sigma, 45 | non_linear_sigma=cfg.asim_linear_non_linear_sigma, 46 | act_enc=cfg.asim_linear_act_enc, 47 | signed_act=cfg.asim_linear_signed_act, 48 | layer=cfg.asim_linear_layer, 49 | hybrid_levels=cfg.asim_cnn_hybrid_levels, 50 | mode=cfg.asim_linear_mode, 51 | trim_noise=cfg.asim_linear_trim_noise, 52 | device=cfg.device) 53 | 54 | 55 | def quant_conv(in_planes, out_planes, stride=1): 56 | return QuantConv2d(in_planes, 57 | out_planes, 58 | kernel_size=3, 59 | stride=stride, 60 | padding=1, 61 | bias=True, 62 | wbit=cfg.quant_conv_wbit, 63 | xbit=cfg.quant_conv_xbit, 64 | signed_act=True, 65 | mode=cfg.quant_conv_mode, 66 | device=cfg.device) 67 | 68 | 69 | if cfg.PACT: 70 | relu = PACT() 71 | else: 72 | relu = nn.ReLU(inplace=True) 73 | 74 | 75 | class VGG(nn.Module): 76 | def __init__(self, features, num_classes=cfg.cls_num, init_weights=True): 77 | super(VGG, self).__init__() 78 | self.features = features 79 | self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) 80 | self.classifier = nn.Sequential( 81 | asim_affine(512 * 7 * 7, 4096), 82 | relu, 83 | nn.Dropout(), 84 | asim_affine(4096, 4096), 85 | relu, 86 | nn.Dropout(), 87 | asim_affine(4096, num_classes) 88 | ) 89 | if init_weights: 90 | self._initialize_weights() 91 | 92 | def forward(self, x): 93 | x = self.features(x) 94 | x = self.avgpool(x) 95 | x = x.view(x.size(0), -1) 96 | x = self.classifier(x) 97 | return x 98 | 99 | def _initialize_weights(self): 100 | for m in self.modules(): 101 | if isinstance(m, nn.Conv2d): 102 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 103 | if m.bias is not None: 104 | nn.init.constant_(m.bias, 0) 105 | elif isinstance(m, ASiMConv2d): 106 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 107 | if m.bias is not None: 108 | nn.init.constant_(m.bias, 0) 109 | elif isinstance(m, QuantConv2d): 110 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 111 | if m.bias is not None: 112 | nn.init.constant_(m.bias, 0) 113 | elif isinstance(m, nn.BatchNorm2d): 114 | nn.init.constant_(m.weight, 1) 115 | nn.init.constant_(m.bias, 0) 116 | elif isinstance(m, nn.Linear): 117 | nn.init.normal_(m.weight, 0, 0.01) 118 | nn.init.constant_(m.bias, 0) 119 | elif isinstance(m, ASiMLinear): 120 | nn.init.normal_(m.weight, 0, 0.01) 121 | nn.init.constant_(m.bias, 0) 122 | 123 | 124 | def make_layers(str, batch_norm=False): 125 | layers = [] 126 | if batch_norm: 127 | layers += [quant_conv(3, 64), 128 | nn.BatchNorm2d(64), 129 | nn.ReLU(inplace=True)] 130 | else: 131 | layers += [quant_conv(3, 64), 132 | nn.ReLU(inplace=True)] 133 | 134 | in_channels = 64 135 | for v in str: 136 | if v == 'M': 137 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 138 | else: 139 | conv2d = asim_conv(in_channels, v) 140 | if batch_norm: 141 | layers += [conv2d, nn.BatchNorm2d(v), relu] 142 | else: 143 | layers += [conv2d, relu] 144 | in_channels = v 145 | return nn.Sequential(*layers) 146 | 147 | 148 | if cfg.large_model: 149 | 150 | str = { 151 | "A": ["M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], 152 | "B": [64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], 153 | "D": [64, "M", 128, 128, "M", 256, 256, 256, "M", 512, 512, 512, "M", 512, 512, 512, "M"], 154 | "E": [64, "M", 128, 128, "M", 256, 256, 256, 256, "M", 512, 512, 512, 512, "M", 512, 512, 512, 512, "M"] 155 | } 156 | else: 157 | str = { 158 | "A": [128, "M", 256, 256, 512, 512, 512, 512, "M"], 159 | "B": [64, 128, 128, "M", 256, 256, 512, 512, 512, 512, "M"], 160 | "D": [64, 128, 128, "M", 256, 256, 256, 512, 512, 512, 512, 512, 512, "M"], 161 | "E": [64, 128, 128, "M", 256, 256, 256, 256, 512, 512, 512, 512, 512, 512, 512, 512, "M"] 162 | } 163 | 164 | 165 | def vgg11_asim(**kwargs): 166 | model = VGG(make_layers(str['A']), **kwargs) 167 | return model 168 | 169 | 170 | def vgg11_bn_asim(**kwargs): 171 | model = VGG(make_layers(str['A'], batch_norm=True), **kwargs) 172 | return model 173 | 174 | 175 | def vgg13_asim(**kwargs): 176 | model = VGG(make_layers(str['B']), **kwargs) 177 | return model 178 | 179 | 180 | def vgg13_bn_asim(**kwargs): 181 | model = VGG(make_layers(str['B'], batch_norm=True), **kwargs) 182 | return model 183 | 184 | 185 | def vgg16_asim(**kwargs): 186 | model = VGG(make_layers(str['D']), **kwargs) 187 | return model 188 | 189 | 190 | def vgg16_bn_asim(**kwargs): 191 | model = VGG(make_layers(str['D'], batch_norm=True), **kwargs) 192 | return model 193 | 194 | 195 | def vgg19_asim(**kwargs): 196 | model = VGG(make_layers(str['E']), **kwargs) 197 | return model 198 | 199 | 200 | def vgg19_bn_asim(**kwargs): 201 | model = VGG(make_layers(str['E'], batch_norm=True), **kwargs) 202 | return model 203 | 204 | 205 | if __name__ == '__main__': 206 | 207 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 208 | 209 | model = vgg16_bn_asim() 210 | model.to(device=device) 211 | 212 | for name, module in model.named_modules(): 213 | print('Layer name: {}, Layer instance: {}'.format(name, module)) 214 | 215 | # Forward 216 | fake_img = torch.randn((1, 3, 28, 28), device=device) 217 | output = model(fake_img) 218 | print(output.shape) 219 | -------------------------------------------------------------------------------- /tools/common_tools.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Some useful functions to be used in this project. 4 | 5 | """ 6 | import logging 7 | import os 8 | import random 9 | import torch 10 | import matplotlib.pyplot as plt 11 | import numpy as np 12 | from datetime import datetime 13 | from PIL import Image 14 | import torchvision.transforms as transforms 15 | 16 | 17 | def setup_seed(seed=6666): 18 | np.random.seed(seed) 19 | random.seed(seed) 20 | torch.manual_seed(seed) # CPU 21 | if torch.cuda.is_available(): 22 | torch.cuda.manual_seed_all(seed) 23 | torch.backends.cudnn.deterministic = True 24 | torch.backends.cudnn.benchmark = True 25 | # Accelerate training if training set does not change drastically. Best config for fixed cudnn (conv, etc.). 26 | 27 | 28 | def inverse_transform(img_, transform_train): 29 | """ 30 | Inverse transform dataset 31 | :param img_: Tensor 32 | :param transform_train: torchvision.transforms 33 | :return: PIL image 34 | """ 35 | if 'Normalize' in str(transform_train): 36 | norm_transform = list(filter(lambda x: isinstance(x, transforms.Normalize), transform_train.transforms)) 37 | mean = torch.tensor(norm_transform[0].mean, dtype=img_.dtype, device=img_.device) 38 | std = torch.tensor(norm_transform[0].std, dtype=img_.dtype, device=img_.device) 39 | img_.mul_(std[:, None, None]).add_(mean[:, None, None]) 40 | 41 | img_ = img_.transpose(0, 2).transpose(0, 1) # C*H*W --> H*W*C 42 | if 'ToTensor' in str(transform_train) or img_.max() < 1: 43 | img_ = img_.detach().numpy() * 255 44 | 45 | if img_.shape[2] == 3: 46 | img_ = Image.fromarray(img_.astype('uint8')).convert('RGB') 47 | elif img_.shape[2] == 1: 48 | img_ = Image.fromarray(img_.astype('uint8').squeeze()) 49 | else: 50 | raise Exception('Invalid img shape, expected 1 or 3 in axis 2, but got {}!'.format(img_.shape[2])) 51 | 52 | return img_ 53 | 54 | 55 | def plot_result(train_x, train_y, valid_x, valid_y, mode, dir): 56 | """ 57 | Plot loss/acc transfer curves for train & valid set. 58 | :param train_x: Epoch 59 | :param train_y: Scalar 60 | :param valid_x: Epoch 61 | :param valid_y: Scalar 62 | :param mode: 'loss' or 'acc' 63 | :param dir: Save figs to which dir 64 | """ 65 | plt.plot(train_x, train_y, label='Train') 66 | plt.plot(valid_x, valid_y, label='Valid') 67 | 68 | plt.ylabel(str(mode)) 69 | plt.xlabel('Epoch') 70 | 71 | location = 'upper right' if mode == 'loss' else 'upper left' 72 | plt.legend(loc=location) 73 | 74 | plt.title('_'.join([mode])) 75 | plt.savefig(os.path.join(dir, mode + '.png')) 76 | plt.close() 77 | 78 | 79 | def show_confmat(confusion_mat, classes, set_name, out_dir, epoch=999, verbose=False, figsize=None, perc=False): 80 | """ 81 | Plot confusion matrix and save the figure 82 | :param confusion_mat: ndarray 83 | :param classes: list or tuple for class name 84 | :param set_name: str, dataset name, train or valid or test 85 | :param out_dir: str, directory to save figure 86 | :param epoch: int, @ which epoch 87 | :param verbose: bool, print information or not 88 | :param figsize: figure size 89 | :param perc: bool, show percentage 90 | :return: 91 | """ 92 | class_num = len(classes) 93 | 94 | # Normalization 95 | confusion_mat_tmp = confusion_mat.copy() 96 | for i in range(len(classes)): 97 | confusion_mat_tmp[i, :] = confusion_mat[i, :] / confusion_mat[i, :].sum() 98 | 99 | # Fig size setup 100 | if class_num < 10: 101 | figsize = 6 102 | elif class_num >= 100: 103 | figsize = 30 104 | else: 105 | figsize = np.linspace(6, 30, 91)[class_num-10] 106 | plt.figure(figsize=(int(figsize), int(figsize*1.3))) 107 | 108 | # Select color 109 | cmap = plt.cm.get_cmap('Greys') # Color label: https://matplotlib.org/2.0.2/examples/color/colormaps_reference.html 110 | plt.imshow(confusion_mat_tmp, cmap=cmap) # Color as percentage 111 | plt.colorbar(fraction=0.03) 112 | 113 | # Plot setup 114 | xlocations = np.array(range(len(classes))) 115 | plt.xticks(xlocations, list(classes), rotation=60) 116 | plt.yticks(xlocations, list(classes)) 117 | plt.xlabel('Predict label') 118 | plt.ylabel('True label') 119 | plt.title('Confusion_Matrix_{}_{}'.format(set_name, epoch)) 120 | 121 | # Abs or % 122 | if perc: 123 | class_per_nums = confusion_mat.sum(axis=0) 124 | conf_mat_perc = confusion_mat / class_per_nums 125 | for i in range(confusion_mat_tmp.shape[0]): 126 | for j in range(confusion_mat_tmp.shape[1]): 127 | plt.text(x=j, y=i, s='{:.0%}' 128 | .format(conf_mat_perc[i, j]), va='center', ha='center', color='red', fontsize=10) 129 | else: 130 | for i in range(confusion_mat_tmp.shape[0]): 131 | for j in range(confusion_mat_tmp.shape[1]): 132 | plt.text(x=j, y=i, s=int(confusion_mat[i, j]), va='center', ha='center', color='red', fontsize=10) 133 | 134 | # Save fig 135 | plt.savefig(os.path.join(out_dir, 'Confusion_Matrix_{}.png'.format(set_name))) 136 | plt.close() 137 | 138 | if verbose: 139 | for i in range(class_num): 140 | print('class: {:<10}, total num: {:<6}, correct num: {:<5} Recall: {:.2%} Precision: {:.2%}' 141 | .format(classes[i], np.sum(confusion_mat[i, :]), confusion_mat[i, i], 142 | confusion_mat[i, i] / (1e-9 + np.sum(confusion_mat[i, :])), 143 | confusion_mat[i, i] / (1e-9 + np.sum(confusion_mat[:, i])))) 144 | 145 | 146 | def check_data_dir(dir): 147 | """ 148 | Check directory. 149 | """ 150 | assert os.path.exists(dir), \ 151 | '\n\nDirectory does not exist. Current variable point to: \n{}\nPlease check relative directory settings or data files'.format( 152 | os.path.abspath(dir)) 153 | 154 | 155 | def make_logger(dir): 156 | """ 157 | Create log file with input dir named with current time as recorder. 158 | """ 159 | now_time = datetime.now() 160 | time_str = datetime.strftime(now_time, '%y-%m-%d_%H-%M') 161 | log_dir = os.path.join(dir, time_str) # Create folder according to config 162 | if not os.path.exists(log_dir): 163 | os.makedirs(log_dir) 164 | # Create logger 165 | path_log = os.path.join(log_dir, 'log.log') 166 | logger = Logger(path_log) 167 | logger = logger.init_logger() 168 | 169 | return logger, log_dir 170 | 171 | 172 | class Logger(): 173 | def __init__(self, path_log): 174 | log_name = os.path.basename(path_log) 175 | self.log_name = log_name if log_name else 'root' 176 | self.out_path = path_log 177 | 178 | log_dir = os.path.dirname(self.out_path) 179 | if not os.path.exists(log_dir): 180 | os.makedirs(log_dir) 181 | 182 | def init_logger(self): 183 | logger = logging.getLogger(self.log_name) 184 | logger.setLevel(level=logging.INFO) 185 | 186 | # config file Handler 187 | file_handler = logging.FileHandler(self.out_path, 'w') 188 | file_handler.setLevel(logging.INFO) 189 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 190 | file_handler.setFormatter(formatter) 191 | 192 | # config console Handler 193 | console_handler = logging.StreamHandler() 194 | console_handler.setLevel(logging.INFO) 195 | # console_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) 196 | 197 | # Add Handler 198 | logger.addHandler(file_handler) 199 | logger.addHandler(console_handler) 200 | 201 | return logger 202 | 203 | 204 | if __name__ == '__main__': # Testbench 205 | setup_seed(6666) 206 | print(np.random.randint(0, 10, 1)) 207 | 208 | logger = Logger('../resnet_cifar100/run/logtest.log') 209 | logger = logger.init_logger() 210 | for i in range(10): 211 | logger.info('test:' + str(i)) 212 | 213 | from main.config import cfg 214 | logger.info(cfg) 215 | -------------------------------------------------------------------------------- /model/resnet_asim.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | ResNet models for the simulation framework. 4 | 5 | """ 6 | 7 | import torch.nn as nn 8 | from module.basic_module import QuantConv2d, PACT 9 | from module.asim_conv import ASiMConv2d 10 | from module.asim_linear import ASiMLinear 11 | from main.config import cfg 12 | import torch 13 | import time 14 | 15 | 16 | def asim_conv3x3(in_planes, out_planes, stride=1): 17 | """3x3 convolution with padding""" 18 | return ASiMConv2d(in_planes, 19 | out_planes, 20 | kernel_size=3, 21 | stride=stride, 22 | padding=1, 23 | bias=False, 24 | wbit=cfg.asim_conv_wbit, 25 | xbit=cfg.asim_conv_xbit, 26 | adc_prec=cfg.asim_conv_adc_prec, 27 | nrow=cfg.asim_conv_nrow, 28 | rand_noise_sigma=cfg.asim_conv_rand_noise_sigma, 29 | non_linear_sigma=cfg.asim_conv_non_linear_sigma, 30 | act_enc=cfg.asim_act_enc, 31 | signed_act=cfg.asim_signed_act, 32 | hybrid_levels=cfg.asim_cnn_hybrid_levels, 33 | mode=cfg.asim_conv_mode, 34 | trim_noise=cfg.asim_conv_trim_noise, 35 | device=cfg.device) 36 | 37 | 38 | def asim_conv1x1(in_planes, out_planes, stride=1): 39 | """1x1 convolution""" 40 | return ASiMConv2d(in_planes, 41 | out_planes, 42 | kernel_size=1, 43 | stride=stride, 44 | bias=False, 45 | wbit=cfg.asim_conv_wbit, 46 | xbit=cfg.asim_conv_xbit, 47 | adc_prec=cfg.asim_conv_adc_prec, 48 | nrow=cfg.asim_conv_nrow, 49 | rand_noise_sigma=cfg.asim_conv_rand_noise_sigma, 50 | non_linear_sigma=cfg.asim_conv_non_linear_sigma, 51 | act_enc=cfg.asim_act_enc, 52 | signed_act=cfg.asim_signed_act, 53 | hybrid_levels=cfg.asim_cnn_hybrid_levels, 54 | mode=cfg.asim_conv_mode, 55 | trim_noise=cfg.asim_conv_trim_noise, 56 | device=cfg.device) 57 | 58 | 59 | def asim_affine(in_features, out_features): 60 | return ASiMLinear(in_features, 61 | out_features, 62 | bias=True, 63 | wbit=cfg.asim_linear_wbit, 64 | xbit=cfg.asim_linear_xbit, 65 | adc_prec=cfg.asim_linear_adc_prec, 66 | nrow=cfg.asim_linear_nrow, 67 | rand_noise_sigma=cfg.asim_linear_rand_noise_sigma, 68 | non_linear_sigma=cfg.asim_linear_non_linear_sigma, 69 | act_enc=cfg.asim_linear_act_enc, 70 | signed_act=cfg.asim_linear_signed_act, 71 | layer=cfg.asim_linear_layer, 72 | hybrid_levels=cfg.asim_cnn_hybrid_levels, 73 | mode=cfg.asim_linear_mode, 74 | trim_noise=cfg.asim_linear_trim_noise, 75 | device=cfg.device) 76 | 77 | 78 | def quant_conv1(in_planes, out_planes): 79 | """3x3 convolution with padding""" 80 | return QuantConv2d(in_planes, 81 | out_planes, 82 | kernel_size=3, 83 | stride=1, 84 | padding=1, 85 | bias=False, 86 | wbit=cfg.quant_conv_wbit, 87 | xbit=cfg.quant_conv_xbit, 88 | signed_act=True, 89 | mode=cfg.quant_conv_mode, 90 | device=cfg.device) 91 | 92 | 93 | def quant_conv2(in_planes, out_planes): 94 | """3x3 convolution with padding""" 95 | return QuantConv2d(in_planes, 96 | out_planes, 97 | kernel_size=7, 98 | stride=2, 99 | padding=3, 100 | bias=False, 101 | wbit=cfg.quant_conv_wbit, 102 | xbit=cfg.quant_conv_xbit, 103 | signed_act=True, 104 | mode=cfg.quant_conv_mode, 105 | device=cfg.device) 106 | 107 | 108 | class BasicBlock(nn.Module): 109 | expansion = 1 110 | 111 | def __init__(self, inplanes, planes, stride=1, downsample=None): 112 | super(BasicBlock, self).__init__() 113 | self.conv1 = asim_conv3x3(inplanes, planes, stride) 114 | self.bn1 = nn.BatchNorm2d(planes) 115 | self.conv2 = asim_conv3x3(planes, planes) 116 | self.bn2 = nn.BatchNorm2d(planes) 117 | self.downsample = downsample 118 | self.stride = stride 119 | if cfg.PACT: 120 | self.relu = PACT() 121 | else: 122 | self.relu = nn.ReLU(inplace=True) 123 | 124 | def forward(self, x): 125 | identity = x 126 | 127 | out = self.conv1(x) 128 | out = self.bn1(out) 129 | out = self.relu(out) 130 | 131 | out = self.conv2(out) 132 | out = self.bn2(out) 133 | 134 | if self.downsample is not None: 135 | identity = self.downsample(x) 136 | 137 | out += identity 138 | out = self.relu(out) 139 | 140 | return out 141 | 142 | 143 | class Bottleneck(nn.Module): 144 | expansion = 4 145 | 146 | def __init__(self, inplanes, planes, stride=1, downsample=None): 147 | super(Bottleneck, self).__init__() 148 | self.conv1 = asim_conv1x1(inplanes, planes) 149 | self.bn1 = nn.BatchNorm2d(planes) 150 | self.conv2 = asim_conv3x3(planes, planes, stride) 151 | self.bn2 = nn.BatchNorm2d(planes) 152 | self.conv3 = asim_conv1x1(planes, planes * self.expansion) 153 | self.bn3 = nn.BatchNorm2d(planes * self.expansion) 154 | self.downsample = downsample 155 | self.stride = stride 156 | if cfg.PACT: 157 | self.relu = PACT() 158 | else: 159 | self.relu = nn.ReLU(inplace=True) 160 | 161 | def forward(self, x): 162 | identity = x 163 | 164 | out = self.conv1(x) 165 | out = self.bn1(out) 166 | out = self.relu(out) 167 | 168 | out = self.conv2(out) 169 | out = self.bn2(out) 170 | out = self.relu(out) 171 | 172 | out = self.conv3(out) 173 | out = self.bn3(out) 174 | 175 | if self.downsample is not None: 176 | identity = self.downsample(x) 177 | 178 | out += identity 179 | out = self.relu(out) 180 | 181 | return out 182 | 183 | 184 | class ResNet(nn.Module): 185 | 186 | def __init__(self, block, layers, num_classes=cfg.cls_num, zero_init_residual=False): 187 | super(ResNet, self).__init__() 188 | self.inplanes = 64 189 | self.conv1 = quant_conv1(3, 64) 190 | self.conv2 = quant_conv2(3, 64) 191 | self.bn1 = nn.BatchNorm2d(64) 192 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 193 | self.layer1 = self._make_layer(block, 64, layers[0]) 194 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 195 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 196 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 197 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 198 | self.fc = asim_affine(512 * block.expansion, num_classes) 199 | if cfg.PACT: 200 | self.relu = PACT() 201 | else: 202 | self.relu = nn.ReLU(inplace=True) 203 | 204 | for m in self.modules(): 205 | if isinstance(m, nn.Conv2d): 206 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 207 | elif isinstance(m, ASiMConv2d): 208 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 209 | elif isinstance(m, QuantConv2d): 210 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 211 | elif isinstance(m, nn.BatchNorm2d): 212 | nn.init.constant_(m.weight, 1) 213 | nn.init.constant_(m.bias, 0) 214 | 215 | if zero_init_residual: 216 | for m in self.modules(): 217 | if isinstance(m, Bottleneck): 218 | nn.init.constant_(m.bn3.weight, 0) 219 | elif isinstance(m, BasicBlock): 220 | nn.init.constant_(m.bn2.weight, 0) 221 | 222 | def _make_layer(self, block, planes, blocks, stride=1): 223 | downsample = None 224 | if stride != 1 or self.inplanes != planes * block.expansion: 225 | downsample = nn.Sequential( 226 | asim_conv1x1(self.inplanes, planes * block.expansion, stride), 227 | nn.BatchNorm2d(planes * block.expansion) 228 | ) 229 | 230 | layers = [] 231 | layers.append(block(self.inplanes, planes, stride, downsample)) 232 | self.inplanes = planes * block.expansion 233 | for _ in range(1, blocks): 234 | layers.append(block(self.inplanes, planes)) 235 | 236 | return nn.Sequential(*layers) 237 | 238 | def forward(self, x): 239 | if cfg.large_model: 240 | x = self.conv2(x) 241 | x = self.bn1(x) 242 | x = self.relu(x) 243 | x = self.maxpool(x) 244 | else: 245 | x = self.conv1(x) 246 | x = self.bn1(x) 247 | x = self.relu(x) 248 | 249 | x = self.layer1(x) 250 | x = self.layer2(x) 251 | x = self.layer3(x) 252 | x = self.layer4(x) 253 | 254 | x = self.avgpool(x) 255 | x = x.view(x.size(0), -1) 256 | x = self.fc(x) 257 | 258 | return x 259 | 260 | 261 | def resnet18_asim(**kwargs): 262 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 263 | 264 | return model 265 | 266 | 267 | def resnet34_asim(**kwargs): 268 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 269 | 270 | return model 271 | 272 | 273 | def resnet50_asim(**kwargs): 274 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 275 | 276 | return model 277 | 278 | 279 | def resnet101_asim(**kwargs): 280 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 281 | 282 | return model 283 | 284 | 285 | def resnet152_asim(**kwargs): 286 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 287 | 288 | return model 289 | 290 | 291 | if __name__ == '__main__': 292 | 293 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 294 | 295 | start = time.time() 296 | 297 | import torch 298 | model = resnet18_asim() 299 | 300 | for name, module in model.named_modules(): 301 | print('Layer name: {}, Layer instance: {}'.format(name, module)) 302 | in_feat_num = model.fc.in_features 303 | model.fc = nn.Linear(in_feat_num, 10) 304 | 305 | model.to(device=device) 306 | 307 | dummy_img = torch.randn(1, 3, 32, 32, device=device) 308 | output = model(dummy_img) 309 | 310 | end = time.time() 311 | print('Execution Time: {}s.'.format(end-start)) 312 | 313 | print(output.shape) 314 | -------------------------------------------------------------------------------- /model/vision_transformer_asim.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Vision Transformer models for the simulation framework. 4 | 5 | """ 6 | 7 | import torch 8 | import torch.nn as nn 9 | from main.config import cfg 10 | from module.basic_module import QuantConv2d 11 | from module.asim_linear import ASiMLinear 12 | from module.asim_multiheadattention import ASiMMultiheadAttention 13 | 14 | 15 | def attn(embed_dim, num_heads, dropout=0.1): 16 | return ASiMMultiheadAttention(embed_dim, 17 | num_heads, 18 | dropout=dropout, 19 | qk_qbit=cfg.asim_vit_attn_qk_qbit, 20 | qk_kbit=cfg.asim_vit_attn_qk_kbit, 21 | av_abit=cfg.asim_vit_attn_av_abit, 22 | av_vbit=cfg.asim_vit_attn_av_vbit, 23 | proj_wbit=cfg.asim_vit_attn_proj_wbit, 24 | proj_xbit=cfg.asim_vit_attn_proj_xbit, 25 | qk_adc_prec=cfg.asim_vit_attn_qk_adc_prec, 26 | av_adc_prec=cfg.asim_vit_attn_av_adc_prec, 27 | proj_adc_prec=cfg.asim_vit_attn_proj_adc_prec, 28 | nrow=cfg.asim_vit_attn_nrow, 29 | qk_rand_noise_sigma=cfg.asim_vit_attn_qk_rand_noise_sigma, 30 | av_rand_noise_sigma=cfg.asim_vit_attn_av_rand_noise_sigma, 31 | proj_rand_noise_sigma=cfg.asim_vit_attn_proj_rand_noise_sigma, 32 | qk_non_linear_sigma=cfg.asim_vit_attn_qk_non_linear_sigma, 33 | av_non_linear_sigma=cfg.asim_vit_attn_av_non_linear_sigma, 34 | proj_non_linear_sigma=cfg.asim_vit_attn_proj_non_linear_sigma, 35 | qk_k_enc=cfg.asim_vit_attn_qk_k_enc, 36 | av_a_enc=cfg.asim_vit_attn_av_a_enc, 37 | proj_act_enc=cfg.asim_vit_attn_proj_act_enc, 38 | hybrid_levels=cfg.asim_vit_hybrid_levels, 39 | mode=cfg.asim_vit_attn_mode, 40 | attn_trim_noise=cfg.asim_vit_attn_attn_trim_noise, 41 | proj_trim_noise=cfg.asim_vit_attn_proj_trim_noise, 42 | device=cfg.device) 43 | 44 | 45 | def mlp(in_features, out_features): 46 | return ASiMLinear(in_features, 47 | out_features, 48 | bias=True, 49 | wbit=cfg.asim_vit_mlp_wbit, 50 | xbit=cfg.asim_vit_mlp_xbit, 51 | adc_prec=cfg.asim_vit_mlp_adc_prec, 52 | nrow=cfg.asim_vit_mlp_nrow, 53 | rand_noise_sigma=cfg.asim_vit_mlp_rand_noise_sigma, 54 | non_linear_sigma=cfg.asim_vit_mlp_non_linear_sigma, 55 | act_enc=cfg.asim_vit_mlp_act_enc, 56 | signed_act=True, 57 | layer='proj', 58 | hybrid_levels=cfg.asim_vit_hybrid_levels, 59 | mode=cfg.asim_vit_mlp_mode, 60 | trim_noise=cfg.asim_vit_mlp_trim_noise, 61 | device=cfg.device) 62 | 63 | 64 | def fc(in_features, out_features): 65 | return ASiMLinear(in_features, 66 | out_features, 67 | bias=True, 68 | wbit=cfg.asim_vit_fc_wbit, 69 | xbit=cfg.asim_vit_fc_xbit, 70 | adc_prec=cfg.asim_vit_fc_adc_prec, 71 | nrow=cfg.asim_vit_fc_nrow, 72 | rand_noise_sigma=cfg.asim_vit_fc_rand_noise_sigma, 73 | non_linear_sigma=cfg.asim_vit_fc_non_linear_sigma, 74 | act_enc=cfg.asim_vit_fc_act_enc, 75 | signed_act=True, 76 | layer='fc', 77 | hybrid_levels=cfg.asim_vit_hybrid_levels, 78 | mode=cfg.asim_vit_fc_mode, 79 | trim_noise=cfg.asim_vit_fc_trim_noise, 80 | device=cfg.device) 81 | 82 | 83 | def quant_conv(in_planes, out_planes, kernel_size, stride): 84 | return QuantConv2d(in_planes, 85 | out_planes, 86 | kernel_size=kernel_size, 87 | stride=stride, 88 | padding=0, 89 | bias=True, 90 | wbit=cfg.asim_vit_quant_conv_wbit, 91 | xbit=cfg.asim_vit_quant_conv_xbit, 92 | signed_act=True, 93 | mode=cfg.asim_vit_quant_conv_mode, 94 | device=cfg.device) 95 | 96 | 97 | class PatchEmbedding(nn.Module): 98 | def __init__(self, image_size, patch_size, in_channels, out_channels): 99 | super().__init__() 100 | self.patch_size = patch_size 101 | self.num_patches = (image_size // patch_size) ** 2 102 | if cfg.vit_quant_conv_proj: 103 | self.conv_proj = quant_conv(in_channels, out_channels, patch_size, patch_size) 104 | else: 105 | self.conv_proj = nn.Conv2d(in_channels, out_channels, kernel_size=patch_size, stride=patch_size) 106 | 107 | def forward(self, x): 108 | x = self.conv_proj(x).flatten(2).transpose(1, 2) 109 | return x 110 | 111 | 112 | class MLPBlock(nn.Module): 113 | def __init__(self, in_dim, mlp_dim=None, out_dim=None, dropout=0.0): 114 | super().__init__() 115 | out_dim = out_dim or in_dim 116 | mlp_dim = mlp_dim or in_dim 117 | self.mlp1 = mlp(in_dim, mlp_dim) 118 | self.act = nn.GELU() 119 | self.mlp2 = mlp(mlp_dim, out_dim) 120 | self.drop = nn.Dropout(dropout) 121 | 122 | def forward(self, x): 123 | x = self.mlp1(x) 124 | x = self.act(x) 125 | x = self.drop(x) 126 | x = self.mlp2(x) 127 | x = self.drop(x) 128 | return x 129 | 130 | 131 | class EncoderBlock(nn.Module): 132 | def __init__(self, num_heads, hidden_dim, mlp_dim, dropout, attention_dropout): 133 | super().__init__() 134 | 135 | self.ln_1 = nn.LayerNorm(hidden_dim) 136 | self.self_attention = attn(hidden_dim, num_heads, dropout=attention_dropout) 137 | self.dropout = nn.Dropout(dropout) 138 | 139 | self.ln_2 = nn.LayerNorm(hidden_dim) 140 | self.mlp = MLPBlock(in_dim=hidden_dim, mlp_dim=mlp_dim, dropout=dropout) 141 | 142 | def forward(self, input): 143 | x = self.ln_1(input) 144 | x = self.self_attention(x) 145 | x = self.dropout(x) 146 | x = x + input 147 | 148 | y = self.ln_2(x) 149 | y = self.mlp(y) 150 | 151 | return x + y 152 | 153 | 154 | class Encoder(nn.Module): 155 | def __init__(self, seq_length, num_layers, num_heads, hidden_dim, mlp_dim, dropout, attention_dropout): 156 | super().__init__() 157 | 158 | self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02)) 159 | self.dropout = nn.Dropout(dropout) 160 | self.layers = nn.Sequential( 161 | *[EncoderBlock(num_heads, hidden_dim, mlp_dim, dropout, attention_dropout) for _ in range(num_layers)] 162 | ) 163 | self.ln = nn.LayerNorm(hidden_dim) 164 | 165 | def forward(self, input): 166 | input = input + self.pos_embedding 167 | return self.ln(self.layers(self.dropout(input))) 168 | 169 | 170 | class VisionTransformer(nn.Module): 171 | def __init__(self, image_size, patch_size, num_layers, num_heads, hidden_dim, mlp_dim, 172 | dropout, attention_dropout, num_classes): 173 | super().__init__() 174 | self.num_classes = num_classes 175 | self.patch_embedding = PatchEmbedding(image_size, patch_size, 3, hidden_dim) 176 | self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_dim)) 177 | seq_length = self.patch_embedding.num_patches + 1 178 | self.encoder = Encoder( 179 | seq_length, 180 | num_layers, 181 | num_heads, 182 | hidden_dim, 183 | mlp_dim, 184 | dropout, 185 | attention_dropout) 186 | self.heads = fc(hidden_dim, num_classes) 187 | 188 | def forward(self, x): 189 | B = x.shape[0] 190 | x = self.patch_embedding(x) 191 | cls_tokens = self.cls_token.expand(B, -1, -1) 192 | x = torch.cat((cls_tokens, x), dim=1) 193 | x = self.encoder(x) 194 | x = x[:, 0] 195 | x = self.heads(x) 196 | 197 | return x 198 | 199 | 200 | def vit_b_16_asim(): 201 | model = VisionTransformer(image_size=cfg.image_size, 202 | patch_size=16, 203 | num_layers=12, 204 | num_heads=12, 205 | hidden_dim=768, 206 | mlp_dim=3072, 207 | dropout=0.1, 208 | attention_dropout=0.1, 209 | num_classes=cfg.cls_num) 210 | return model 211 | 212 | 213 | def vit_b_32_asim(): 214 | model = VisionTransformer(image_size=cfg.image_size, 215 | patch_size=32, 216 | num_layers=12, 217 | num_heads=12, 218 | hidden_dim=768, 219 | mlp_dim=3072, 220 | dropout=0.1, 221 | attention_dropout=0.1, 222 | num_classes=cfg.cls_num) 223 | return model 224 | 225 | 226 | def vit_l_16_asim(): 227 | model = VisionTransformer(image_size=cfg.image_size, 228 | patch_size=16, 229 | num_layers=24, 230 | num_heads=16, 231 | hidden_dim=1024, 232 | mlp_dim=4096, 233 | dropout=0.1, 234 | attention_dropout=0.1, 235 | num_classes=cfg.cls_num) 236 | return model 237 | 238 | 239 | def vit_l_32_asim(): 240 | model = VisionTransformer(image_size=cfg.image_size, 241 | patch_size=32, 242 | num_layers=24, 243 | num_heads=16, 244 | hidden_dim=1024, 245 | mlp_dim=4096, 246 | dropout=0.1, 247 | attention_dropout=0.1, 248 | num_classes=cfg.cls_num) 249 | return model 250 | 251 | 252 | def vit_h_14_asim(): 253 | model = VisionTransformer(image_size=cfg.image_size, 254 | patch_size=14, 255 | num_layers=32, 256 | num_heads=16, 257 | hidden_dim=1280, 258 | mlp_dim=5120, 259 | dropout=0.1, 260 | attention_dropout=0.1, 261 | num_classes=cfg.cls_num) 262 | return model 263 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /examples/train/resnet18_cifar10_qat/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | # cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | # cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | origin_size = 32 42 | input_size = 32 43 | # ImageNet and ViT 44 | # origin_size = 256 45 | # input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 10 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Train' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Train' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 8 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 8 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Train' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Train' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Train' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | # cfg.train_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | # cfg.valid_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | # cfg.test_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | cfg.pretrain_model_path = None # For training, give a None. 178 | # cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\vitb32_imagenet_w8a8_trim_40.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 128 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /examples/train/resnet18_cifar10_nat/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | # cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | # cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | origin_size = 32 42 | input_size = 32 43 | # ImageNet and ViT 44 | # origin_size = 256 45 | # input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 10 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Train' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 50.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Train' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 50.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 8 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 8 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Train' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Train' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Train' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | # cfg.train_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | # cfg.valid_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | # cfg.test_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\resnet18_cifar10_w8a8_pact_trim_100.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 128 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 80 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.001 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /examples/simulation/resnet18_imagenet/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | # cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | # cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | # origin_size = 32 42 | # input_size = 32 43 | # ImageNet and ViT 44 | origin_size = 256 45 | input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 1000 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Simulation' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Simulation' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 8 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 8 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Train' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Train' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Train' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | # cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | # cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | # cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | cfg.train_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | cfg.test_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\resnet18_imagenet_w8a8_pact_trim_50.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 128 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /examples/simulation/vitb32_cifar10/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | # cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | # cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | # origin_size = 32 42 | # input_size = 32 43 | # ImageNet and ViT 44 | origin_size = 256 45 | input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 10 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Simulation' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Simulation' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 9 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 9 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Simulation' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Simulation' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Simulation' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | # cfg.train_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | # cfg.valid_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | # cfg.test_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\vitb32_cifar10_w8a8_trim_60.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 64 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /examples/simulation/vitb32_imagenet/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | # cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | # cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | # origin_size = 32 42 | # input_size = 32 43 | # ImageNet and ViT 44 | origin_size = 256 45 | input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 1000 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Simulation' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Simulation' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 9 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 9 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Simulation' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Simulation' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Simulation' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | # cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | # cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | # cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | cfg.train_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | cfg.test_dir = r'D:\Machine_Learning\Dataset\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\vitb32_imagenet_w8a8_trim_40.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 64 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /examples/simulation/resnet18_cifar10/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | # cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | # cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | origin_size = 32 42 | input_size = 32 43 | # ImageNet and ViT 44 | # origin_size = 256 45 | # input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 10 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Simulation' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Simulation' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 8 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 8 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Train' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Train' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Train' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | cfg.train_dir = r'D:\Machine_Learning\Dataset\CIFAR10\train' 165 | cfg.valid_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 166 | cfg.test_dir = r'D:\Machine_Learning\Dataset\CIFAR10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'D:\Machine_Learning\ASiM_Models\resnet18_cifar10_w8a8_pact_trim_100.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 512 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /main/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | Parameter management config for training/simulation. 4 | 5 | """ 6 | 7 | import matplotlib.pyplot as plt 8 | import torchvision.transforms as transforms 9 | from easydict import EasyDict 10 | 11 | # Obtain value via .key like key-value pair 12 | cfg = EasyDict() 13 | 14 | # <|Specify a task|> 15 | cfg.task = 'CIFAR-10' 16 | # cfg.task = 'CIFAR-100' 17 | # cfg.task = 'ImageNet' 18 | 19 | # <|Specify a model|> 20 | cfg.model_name = 'ResNet-18' 21 | # cfg.model_name = 'ResNet-34' 22 | # cfg.model_name = 'ResNet-50' 23 | # cfg.model_name = 'ResNet-101' 24 | # cfg.model_name = 'ResNet-152' 25 | # cfg.model_name = 'VGG11' 26 | # cfg.model_name = 'VGG13' 27 | # cfg.model_name = 'VGG16' 28 | # cfg.model_name = 'VGG19' 29 | # cfg.model_name = 'VGG11-BN' 30 | # cfg.model_name = 'VGG13-BN' 31 | # cfg.model_name = 'VGG16-BN' 32 | # cfg.model_name = 'VGG19-BN' 33 | # cfg.model_name = 'ViT-B-16' 34 | # cfg.model_name = 'ViT-B-32' 35 | # cfg.model_name = 'ViT-L-16' 36 | # cfg.model_name = 'ViT-L-32' 37 | # cfg.model_name = 'ViT-H-14' 38 | 39 | # <|Specify the input image size|> 40 | # CIFAR-10, CIFAR-100 41 | origin_size = 32 42 | input_size = 32 43 | # ImageNet and ViT 44 | # origin_size = 256 45 | # input_size = 224 46 | 47 | # <|Specify the running device: 'cpu', 'cuda', 'mps', etc.|> 48 | cfg.device = 'cuda' 49 | 50 | # <|Specify the classifier dimension|> 51 | cfg.cls_num = 10 52 | 53 | ############################################################# 54 | ################### ################### 55 | ############################################################# 56 | 57 | # <|Specify using PACT or ReLU|> 58 | cfg.PACT = True 59 | # <|Specify Hybrid Computing Settings|> 60 | cfg.asim_cnn_hybrid_levels = None # Hybrid computing boundary level 61 | 62 | # <|Specify the settings for ASiMConv2d layers|> 63 | cfg.asim_conv_mode = 'Simulation' # Mode: Train/Inference/Simulation 64 | cfg.asim_conv_wbit = 8 # Weight bit 65 | cfg.asim_conv_xbit = 8 # Activation bit 66 | cfg.asim_conv_adc_prec = 8 # ADC precision 67 | cfg.asim_conv_nrow = 255 # CiM column length (number of row-parallelism) 68 | cfg.asim_conv_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 69 | cfg.asim_conv_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 70 | cfg.asim_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 71 | cfg.asim_signed_act = False # Signed or unsigned activation 72 | cfg.asim_conv_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 73 | 74 | # <|Specify the settings for ASiMLinear layers|> 75 | cfg.asim_linear_mode = 'Simulation' # Mode: Train/Inference/Simulation 76 | cfg.asim_linear_wbit = 8 # Weight bit 77 | cfg.asim_linear_xbit = 8 # Activation bit 78 | cfg.asim_linear_adc_prec = 8 # ADC precision 79 | cfg.asim_linear_nrow = 255 # CiM column length (number of row-parallelism) 80 | cfg.asim_linear_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 81 | cfg.asim_linear_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 82 | cfg.asim_linear_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 83 | cfg.asim_linear_signed_act = False # Signed or unsigned activation 84 | cfg.asim_linear_layer = 'fc' # Layer type ('fc' (fully-connected) or 'proj' (projection)) 85 | cfg.asim_linear_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 86 | 87 | # <|Specify the settings for QuantConv2d layers|> 88 | cfg.quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 89 | cfg.quant_conv_wbit = 8 # Weight bit 90 | cfg.quant_conv_xbit = 8 # Activation bit 91 | 92 | ############################################################# 93 | ############################################################# 94 | ############################################################# 95 | 96 | ############################################################# 97 | ################### ################### 98 | ############################################################# 99 | 100 | # <|Input image size|> 101 | cfg.image_size = input_size 102 | # <|If use QuantConv2d to extract token embedding|> 103 | cfg.vit_quant_conv_proj = False 104 | # <|Specify Hybrid Computing Settings|> 105 | cfg.asim_vit_hybrid_levels = None # Hybrid computing boundary level 106 | 107 | # <|Specify the settings for ASiMMultiheadAttention layers|> 108 | cfg.asim_vit_attn_qk_qbit = 8 # Query bit in QK computation 109 | cfg.asim_vit_attn_qk_kbit = 8 # Key bit in QK computation 110 | cfg.asim_vit_attn_av_abit = 8 # Score bit in AV computation 111 | cfg.asim_vit_attn_av_vbit = 8 # Value bit in AV computation 112 | cfg.asim_vit_attn_proj_wbit = 8 # Weight bit in WX/WO projection 113 | cfg.asim_vit_attn_proj_xbit = 8 # Activation bit in WX/WO projection 114 | cfg.asim_vit_attn_qk_adc_prec = 8 # ADC precision in QK computation 115 | cfg.asim_vit_attn_av_adc_prec = 8 # ADC precision in AV computation 116 | cfg.asim_vit_attn_proj_adc_prec = 8 # ADC precision in WX/WO projection 117 | cfg.asim_vit_attn_nrow = 255 # CiM column length (number of row-parallelism) 118 | cfg.asim_vit_attn_qk_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (QK computation) 119 | cfg.asim_vit_attn_av_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (AV computation) 120 | cfg.asim_vit_attn_proj_rand_noise_sigma = 0.0 # Standard deviation of random noise in % (WQ/WK/WV/WO computation) 121 | cfg.asim_vit_attn_qk_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (QK computation) 122 | cfg.asim_vit_attn_av_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (AV computation) 123 | cfg.asim_vit_attn_proj_non_linear_sigma = 0.0 # Standard deviation of non-linear in % (WQ/WK/WV/WO computation) 124 | cfg.asim_vit_attn_qk_k_enc = None # Key encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 125 | cfg.asim_vit_attn_av_a_enc = None # Score encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 126 | cfg.asim_vit_attn_proj_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 127 | cfg.asim_vit_attn_mode = 'Train' # Mode: Train/Inference/Simulation 128 | cfg.asim_vit_attn_attn_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Attention) 129 | cfg.asim_vit_attn_proj_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy (Projection) 130 | 131 | # <|Specify the settings for MLP layers|> 132 | cfg.asim_vit_mlp_wbit = 8 # Weight bit 133 | cfg.asim_vit_mlp_xbit = 8 # Activation bit 134 | cfg.asim_vit_mlp_adc_prec = 8 # ADC precision 135 | cfg.asim_vit_mlp_nrow = 255 # CiM column length (number of row-parallelism) 136 | cfg.asim_vit_mlp_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 137 | cfg.asim_vit_mlp_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 138 | cfg.asim_vit_mlp_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 139 | cfg.asim_vit_mlp_mode = 'Train' # Mode: Train/Inference/Simulation 140 | cfg.asim_vit_mlp_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 141 | 142 | # <|Specify the settings for FC layer|> 143 | cfg.asim_vit_fc_wbit = 8 # Weight bit 144 | cfg.asim_vit_fc_xbit = 8 # Activation bit 145 | cfg.asim_vit_fc_adc_prec = 8 # ADC precision 146 | cfg.asim_vit_fc_nrow = 255 # CiM column length (number of row-parallelism) 147 | cfg.asim_vit_fc_rand_noise_sigma = 0.0 # Standard deviation of random noise in % 148 | cfg.asim_vit_fc_non_linear_sigma = 0.0 # Standard deviation of non-linear in % 149 | cfg.asim_vit_fc_act_enc = None # Activation encoding bit for bit-parallel computing (Give None for Bit-Serial Simulation) 150 | cfg.asim_vit_fc_mode = 'Train' # Mode: Train/Inference/Simulation 151 | cfg.asim_vit_fc_trim_noise = 0.0 # Noise intensity for noise-aware training or evaluate inference accuracy 152 | 153 | # <|Specify the settings for QuantConv2d layers|> 154 | cfg.asim_vit_quant_conv_wbit = 8 # Weight bit 155 | cfg.asim_vit_quant_conv_xbit = 8 # Activation bit 156 | cfg.asim_vit_quant_conv_mode = 'Train' # Mode: Train/Inference/Simulation 157 | 158 | ############################################################# 159 | ############################################################# 160 | ############################################################# 161 | 162 | # <|Specify the directory for training/validation/test(simulation) dataset|> 163 | # CIFAR-10 164 | cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar10\archive\cifar10\train' 165 | cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar10\archive\cifar10\test' 166 | cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar10\archive\cifar10\test' 167 | # CIFAR-100 168 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\train_images' 169 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 170 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\Cifar100\test_images' 171 | # ImageNet 172 | # cfg.train_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\train' 173 | # cfg.valid_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 174 | # cfg.test_dir = r'E:\Machine_Learning\Dataset\Processing\ImageNet\imagenet-object-localization-challenge\ILSVRC\Data\CLS-LOC\valid' 175 | 176 | # <|Specify the pretrained model for noise-aware training or test (simulation)|> 177 | # cfg.pretrain_model_path = None # For training, give a None. 178 | cfg.pretrain_model_path = r'E:\Machine_Learning\ASiM_Model_Lib\resnet18_cifar10_w8a8_pact_trim_100.pkl' 179 | 180 | # <|Specify the batch size and workers for training|> 181 | cfg.train_bs = 128 182 | cfg.valid_bs = 128 183 | cfg.train_workers = 8 184 | 185 | # <|Specify the batch size and workers for test (simulation)|> 186 | cfg.test_bs = 128 187 | cfg.test_workers = 8 188 | 189 | ############################################################### 190 | ##################### ##################### 191 | ############################################################### 192 | 193 | # Settings for initial training 194 | cfg.max_epoch = 200 # Training = 200, Fine-Tune = 80 195 | cfg.lr_init = 0.01 # SGD Training = 0.01, SGD Fine-Tune = 0.001, Adam Training = 1e-4, Adam Fine-Tune = 1e-5 196 | cfg.milestones = [int(cfg.max_epoch * 0.6), int(cfg.max_epoch * 0.9)] 197 | cfg.mixup = False # Use mixup or not. 198 | cfg.label_smooth = False # Use label smoothing or not. 199 | 200 | ############################################################### 201 | ############################################################### 202 | ############################################################### 203 | 204 | 205 | ##################################################################################################################### 206 | # No need to modify the following params in general 207 | ##################################################################################################################### 208 | 209 | # <|Specify the optimizer parameters for training/noise-aware training|> 210 | cfg.momentum = 0.9 211 | cfg.weight_decay = 1e-4 212 | cfg.factor = 0.1 # gamma (lr decay param) when using MultiStepLR 213 | cfg.log_interval = 10 # Training log interval 214 | cfg.mixup_alpha = 1.0 # Mixup parameter of beta distribution. 215 | cfg.label_smooth_eps = 0.01 # Label smoothing eps parameter. 216 | 217 | assert cfg.task in ['CIFAR-10', 'CIFAR-100', 'ImageNet'], "Invalid task specified. Choose from 'CIFAR-10', 'CIFAR-100', or 'ImageNet'." 218 | if cfg.task == 'CIFAR-10': 219 | norm_mean = [0.4914, 0.4822, 0.4465] # CIFAR-10 220 | norm_std = [0.2470, 0.2435, 0.2616] # CIFAR-10 221 | if cfg.task == 'CIFAR-100': 222 | norm_mean = [0.5071, 0.4867, 0.4408] # CIFAR-100 223 | norm_std = [0.2675, 0.2565, 0.2761] # CIFAR-100 224 | if cfg.task == 'ImageNet': 225 | norm_mean = [0.485, 0.456, 0.406] # ImageNet 226 | norm_std = [0.229, 0.224, 0.225] # ImageNet 227 | 228 | cfg.transforms_train = transforms.Compose([ 229 | transforms.RandomChoice( 230 | [ 231 | transforms.ColorJitter(brightness=0.5), 232 | transforms.ColorJitter(contrast=0.5), 233 | transforms.ColorJitter(saturation=0.5), 234 | transforms.ColorJitter(hue=0.3) 235 | ] 236 | ), 237 | transforms.Resize((origin_size, origin_size)), 238 | transforms.CenterCrop(origin_size), 239 | transforms.RandomCrop(input_size, padding=4), 240 | transforms.RandomHorizontalFlip(), 241 | transforms.RandomRotation(15), 242 | transforms.ToTensor(), 243 | transforms.Normalize(norm_mean, norm_std) 244 | ]) 245 | 246 | cfg.transforms_valid = transforms.Compose([ 247 | transforms.Resize((input_size, input_size)), 248 | transforms.ToTensor(), 249 | transforms.Normalize(norm_mean, norm_std) 250 | ]) 251 | 252 | cfg.transforms_test = transforms.Compose([ 253 | transforms.Resize((input_size, input_size)), 254 | transforms.ToTensor(), 255 | transforms.Normalize(norm_mean, norm_std) 256 | ]) 257 | 258 | if input_size > 100: 259 | cfg.large_model = True 260 | else: 261 | cfg.large_model = False 262 | 263 | if __name__ == '__main__': # Testbench 264 | 265 | from dataset.cifar100 import Cifar100Dataset 266 | from torch.utils.data import DataLoader 267 | from tools.common_tools import inverse_transform 268 | train_data = Cifar100Dataset(root_dir=cfg.train_dir, transform=cfg.transforms_train) 269 | train_loader = DataLoader(dataset=train_data, batch_size=cfg.train_bs, shuffle=True) 270 | 271 | for epoch in range(cfg.max_epoch): 272 | for i, data in enumerate(train_loader): 273 | 274 | inputs, labels, dir = data # B C H W 275 | 276 | img_tensor = inputs[0, ...] # C H W 277 | img = inverse_transform(img_tensor, cfg.transforms_train) 278 | plt.imshow(img) 279 | plt.show() 280 | plt.pause(0.5) 281 | plt.close() 282 | -------------------------------------------------------------------------------- /module/basic_module.py: -------------------------------------------------------------------------------- 1 | """ 2 | Basic classes and functions for models. 3 | """ 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from tools.common_tools import setup_seed 9 | 10 | 11 | class Floor(torch.autograd.Function): 12 | """ 13 | We can implement our own custom autograd Functions by subclassing 14 | torch.autograd.Function and implementing the forward and backward passes 15 | which operate on Tensors. 16 | """ 17 | 18 | @staticmethod 19 | def forward(ctx, input): 20 | """ 21 | In the forward pass we receive a Tensor containing the input and return 22 | a Tensor containing the output. ctx is a context object that can be used 23 | to stash information for backward computation. You can cache arbitrary 24 | objects for use in the backward pass using the ctx.save_for_backward method. 25 | """ 26 | return torch.floor(input) 27 | 28 | @staticmethod 29 | def backward(ctx, grad_output): 30 | """ 31 | In the backward pass we receive a Tensor containing the gradient of the loss 32 | with respect to the output, and we need to compute the gradient of the loss 33 | with respect to the input. 34 | The backward behavior of the floor function is defined as the identity function. 35 | """ 36 | grad_input = grad_output.clone() 37 | return grad_input 38 | 39 | 40 | class Round(torch.autograd.Function): 41 | """ 42 | We can implement our own custom autograd Functions by subclassing 43 | torch.autograd.Function and implementing the forward and backward passes 44 | which operate on Tensors. 45 | """ 46 | 47 | @staticmethod 48 | def forward(ctx, input): 49 | """ 50 | In the forward pass we receive a Tensor containing the input and return 51 | a Tensor containing the output. ctx is a context object that can be used 52 | to stash information for backward computation. You can cache arbitrary 53 | objects for use in the backward pass using the ctx.save_for_backward method. 54 | """ 55 | return torch.round(input) 56 | 57 | @staticmethod 58 | def backward(ctx, grad_output): 59 | """ 60 | In the backward pass we receive a Tensor containing the gradient of the loss 61 | with respect to the output, and we need to compute the gradient of the loss 62 | with respect to the input. 63 | The backward behavior of the round function is defined as the identity function. 64 | """ 65 | grad_input = grad_output.clone() 66 | return grad_input 67 | 68 | 69 | class Clamp(torch.autograd.Function): 70 | """ 71 | We can implement our own custom autograd Functions by subclassing 72 | torch.autograd.Function and implementing the forward and backward passes 73 | which operate on Tensors. 74 | """ 75 | 76 | @staticmethod 77 | def forward(ctx, input, min, max): 78 | """ 79 | In the forward pass we receive a Tensor containing the input and return 80 | a Tensor containing the output. ctx is a context object that can be used 81 | to stash information for backward computation. You can cache arbitrary 82 | objects for use in the backward pass using the ctx.save_for_backward method. 83 | """ 84 | return torch.clamp(input, min, max) 85 | 86 | @staticmethod 87 | def backward(ctx, grad_output): 88 | """ 89 | In the backward pass we receive a Tensor containing the gradient of the loss 90 | with respect to the output, and we need to compute the gradient of the loss 91 | with respect to the input. 92 | The backward behavior of the clamp function is defined as the identity function. 93 | """ 94 | grad_input = grad_output.clone() 95 | return grad_input, None, None 96 | 97 | 98 | class PACTFunction(torch.autograd.Function): 99 | """ 100 | Gradient computation function for PACT. 101 | """ 102 | @staticmethod 103 | def forward(ctx, input, alpha): 104 | ctx.save_for_backward(input, alpha) 105 | return torch.clamp(input, min=0, max=alpha.item()) 106 | 107 | @staticmethod 108 | def backward(ctx, grad_output): 109 | input, alpha = ctx.saved_tensors 110 | grad_input = grad_output.clone() 111 | grad_alpha = grad_output.clone() 112 | 113 | grad_input[input < 0] = 0 114 | grad_input[input > alpha] = 0 115 | 116 | grad_alpha = grad_alpha * ((input > alpha).float()) 117 | 118 | return grad_input, grad_alpha.sum() 119 | 120 | 121 | class PACT(nn.Module): 122 | """ 123 | PACT ReLU Function 124 | """ 125 | def __init__(self, alpha=6.0): 126 | super(PACT, self).__init__() 127 | self.alpha = nn.Parameter(torch.tensor(alpha)) 128 | 129 | def forward(self, x): 130 | return PACTFunction.apply(x, self.alpha) 131 | 132 | 133 | class QuantConv2d(nn.Conv2d): 134 | """ 135 | Quantized Conv2d Module. 136 | in_channels: input channel 137 | out_channels: output channel 138 | kernel_size: kernel size 139 | stride: stride 140 | padding: padding 141 | dilation: dilation 142 | groups: groups 143 | bias: bias 144 | padding_mode: padding mode 145 | wbit: bit width of weight 146 | xbit: bit width of input activation 147 | signed_act: signed or unsigned activation 148 | mode: operation mode, e.g.: 'Train' or 'Inference' 149 | device: 'cpu' or 'cuda' or 'mps' depend on your device 150 | """ 151 | def __init__(self, 152 | in_channels, 153 | out_channels, 154 | kernel_size, 155 | stride=1, 156 | padding=0, 157 | dilation=1, 158 | groups=1, 159 | bias=True, 160 | padding_mode='zeros', 161 | wbit=8, 162 | xbit=8, 163 | signed_act=False, 164 | mode='Train', 165 | device='cpu'): 166 | super(QuantConv2d, self).__init__(in_channels, 167 | out_channels, 168 | kernel_size, 169 | stride, 170 | padding, 171 | dilation, 172 | groups, 173 | bias, 174 | padding_mode) 175 | assert mode in ['Train', 'Inference'], "Invalid mode specified. Choose from 'Train', 'Inference'." 176 | self.wbit = wbit 177 | self.xbit = xbit 178 | self.signed_act = signed_act 179 | self.kernel_size_param = kernel_size 180 | self.in_channels_param = in_channels 181 | self.padding_param = padding 182 | self.stride_param = stride 183 | self.bias_param = bias 184 | self.epsilon = 1e-7 185 | self.mode = mode 186 | self.device = device 187 | 188 | def _quantize_weight_train(self, input): 189 | """ 190 | Quantize weight tensor in 'Train' mode. 191 | input: input weight tensor 192 | return: fake quantized weight tensor for training 193 | """ 194 | assert self.wbit > 1, "Bit width must be greater than 1." 195 | sign = torch.sign(input).detach() 196 | input = torch.abs(input) 197 | scaling = torch.max(input).detach() + self.epsilon 198 | input = Clamp.apply(input / scaling, 0.0, 1.0) # range from 0~1 199 | input = Round.apply(input * (2.0 ** (self.wbit - 1.0) - 1.0)) / (2.0 ** (self.wbit - 1.0) - 1.0) # INT Quantization 200 | return input * scaling * sign 201 | 202 | def _quantize_signed_feature_train(self, input): 203 | """ 204 | Quantize signed input activation tensor in 'Train' mode. 205 | input: input activation tensor 206 | return: fake quantized input activation tensor for training 207 | """ 208 | assert self.xbit > 1, "Bit width must be greater than 1." 209 | sign = torch.sign(input).detach() 210 | input = torch.abs(input) 211 | scaling = torch.max(input).detach() + self.epsilon 212 | input = Clamp.apply(input / scaling, 0.0, 1.0) # range from 0~1 213 | input = Round.apply(input * (2.0 ** (self.xbit - 1.0) - 1.0)) / (2.0 ** (self.xbit - 1.0) - 1.0) # INT Quantization 214 | return input * scaling * sign 215 | 216 | def _quantize_unsigned_feature_train(self, input): 217 | """ 218 | Quantize unsigned input activation tensor in 'Train' mode. 219 | input: input activation tensor 220 | return: fake quantized input activation tensor for training 221 | """ 222 | assert self.xbit > 1, "Bit width must be greater than 1." 223 | x_scale = torch.max(input).detach() + self.epsilon 224 | input = Clamp.apply(input / x_scale, 0.0, 1.0) # range from 0~1 225 | input = Round.apply(input * (2.0 ** self.xbit - 1.0)) / (2.0 ** self.xbit - 1.0) # UINT Quantization 226 | return input * x_scale 227 | 228 | def _quantize_weight_infer(self, input): 229 | """ 230 | Quantize weight tensor in 'Inference' mode. 231 | input: input weight tensor 232 | return: quantized weight with INT format for inference and scale factor for dequantization 233 | """ 234 | assert self.wbit > 1, "Bit width must be greater than 1." 235 | sign = torch.sign(input).detach() 236 | input = torch.abs(input) 237 | scaling = torch.max(input).detach() + self.epsilon 238 | input = Clamp.apply(input / scaling, 0.0, 1.0) # range from 0~1 239 | input = Round.apply(input * (2.0 ** (self.wbit - 1.0) - 1.0)) # INT Quantization 240 | return input * sign, scaling 241 | 242 | def _quantize_signed_feature_infer(self, input): 243 | """ 244 | Quantize signed input activation tensor in 'Inference' mode. 245 | input: input activation tensor 246 | return: quantized activation with INT format for inference and scale factor for dequantization 247 | """ 248 | assert self.xbit > 1, "Bit width must be greater than 1." 249 | sign = torch.sign(input).detach() 250 | input = torch.abs(input) 251 | scaling = torch.max(input).detach() + self.epsilon 252 | input = Clamp.apply(input / scaling, 0.0, 1.0) # range from 0~1 253 | input = Round.apply(input * (2.0 ** (self.xbit - 1.0) - 1.0)) # INT Quantization 254 | return input * sign, scaling 255 | 256 | def _quantize_unsigned_feature_infer(self, input): 257 | """ 258 | Quantize unsigned input activation tensor in 'Inference' mode. 259 | input: input activation tensor 260 | return: quantized activation with UINT format for inference and scale factor for dequantization 261 | """ 262 | assert self.xbit > 1, "Bit width must be greater than 1." 263 | x_scale = torch.max(input).detach() + self.epsilon 264 | input = Clamp.apply(input / x_scale, 0.0, 1.0) # range from 0~1 265 | input = Round.apply(input * (2.0 ** self.xbit - 1.0)) # UINT Quantization 266 | return input, x_scale 267 | 268 | def forward(self, input): 269 | """ 270 | Forward call of QuantConv2d with selective operation mode. 271 | input: input activation tensor 272 | return: output activation tensor 273 | """ 274 | if self.mode == 'Train': # Training mode 275 | if self.signed_act: 276 | x_quant = self._quantize_signed_feature_train(input) 277 | else: 278 | x_quant = self._quantize_unsigned_feature_train(input) 279 | w_quant = self._quantize_weight_train(self.weight) 280 | output = F.conv2d(x_quant, 281 | w_quant, 282 | self.bias, 283 | self.stride, 284 | self.padding, 285 | self.dilation, 286 | self.groups) 287 | 288 | if self.mode == 'Inference': # Inference mode that mimic quantized output 289 | if self.signed_act: 290 | x_quant, x_scale = self._quantize_signed_feature_infer(input) 291 | else: 292 | x_quant, x_scale = self._quantize_unsigned_feature_infer(input) 293 | w_quant, w_scale = self._quantize_weight_infer(self.weight) 294 | output = F.conv2d(x_quant, 295 | w_quant, 296 | bias=self.bias, 297 | stride=self.stride, 298 | padding=self.padding, 299 | dilation=self.dilation, 300 | groups=self.groups) 301 | # De-quantization to FP32 302 | if self.signed_act: 303 | output = output * (w_scale * x_scale) / ((2.0 ** (self.wbit - 1.0) - 1.0) * (2.0 ** (self.xbit - 1.0) - 1.0)) 304 | else: 305 | output = output * (w_scale * x_scale) / ((2.0 ** (self.wbit - 1.0) - 1.0) * (2.0 ** self.xbit - 1.0)) 306 | # Compensation bias if needed 307 | if self.bias_param: 308 | x_height = output.shape[-2] 309 | x_width = output.shape[-1] 310 | bias_infer = self.bias.repeat(x_height, x_width, 1).permute(2, 0, 1) 311 | output += bias_infer 312 | 313 | return output 314 | 315 | 316 | if __name__ == '__main__': # testbench 317 | setup_seed(6666) 318 | device = 'cpu' 319 | 320 | # fake_conv_img = torch.abs(torch.randn((1, 64, 32, 32), device=device)) 321 | fake_conv_img = torch.randn((1, 64, 32, 32), device=device) 322 | fake_conv_weight = torch.randn((1, 64, 3, 3), device=device) 323 | 324 | model_conv_train = QuantConv2d(64, 325 | 1, 326 | kernel_size=3, 327 | stride=1, 328 | padding=0, 329 | bias=False, 330 | wbit=4, 331 | xbit=4, 332 | mode='Train', 333 | device='cpu') 334 | 335 | model_conv_infer = QuantConv2d(64, 336 | 1, 337 | kernel_size=3, 338 | stride=1, 339 | padding=0, 340 | bias=False, 341 | wbit=4, 342 | xbit=4, 343 | mode='Inference', 344 | device='cpu') 345 | 346 | model_conv_train._parameters['weight'] = fake_conv_weight 347 | model_conv_infer._parameters['weight'] = fake_conv_weight 348 | 349 | output_conv_train = model_conv_train(fake_conv_img) 350 | output_conv_infer = model_conv_infer(fake_conv_img) 351 | 352 | train_to_infer_conv_error = output_conv_infer - output_conv_train 353 | train_to_infer_conv_error_perc = train_to_infer_conv_error / output_conv_train 354 | 355 | print('Conv Layer: Train to Inference Error = {}.'.format(train_to_infer_conv_error_perc)) 356 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ASiM 2 | 3 | ASiM is an inference simulation framework for SRAM-based ACiM design developed by Keio CSG. We hope it would help space exploration of ACiM circuit design and fast inference evaluation on a prototype chip. For any discussion, please feel free to comment on the "Issues" tab, thanks. 4 | 5 | ## 1. Overview 6 | 7 | ### 1.1. Introduction to the Directory 8 | 9 | * `./bins`: Some useful scripts to process datasets and model weights. 10 | * `./dataset`: Datasets for CiM benchmarks. 11 | * `./model`: DNN models for simulations. 12 | * `./module`: ASiM modules for Bit-wise simulation. 13 | * `./tools`: Some useful tools. 14 | * `./main`: Main directory for simulations. 15 | 16 | ### 1.2. Quick Start 17 | 18 | All you need to do is to modify the parameter/argument settings `config.py` in the main folder. Then run `src_train.py` for model training and `src_simulation.py` for simulations. 19 | 20 | When you run `src_train.py`, the training log and model weight will be generated in `main/run`. 21 | 22 | ## 2. Environment 23 | 24 | We handle this project with the following environment: 25 | 26 | ``` 27 | python==3.11.7 28 | torch==2.3.0 29 | numpy==1.24.3 30 | matplotlib==3.8.0 31 | easydict==1.13 32 | ``` 33 | 34 | ## 3. Simulation Time 35 | 36 | A personal GPU will be good to handle most conditions. Simulation can be run with CPU for ResNet-18/CIFAR-10 level task. For ViT simulations on ImageNet, we suggest to run on a GPU server. 37 | 38 | ASiM can be run with various devices and an estimation of running time (w/ full validation set) is listed as follows: 39 | 40 | | Environment | Batch Size | Dataset | Model | Duration (s) | 41 | |-----------------------|------------|----------|------------------|--------------| 42 | | i9-13980HX (64G DDR5) | 128 | CIFAR-10 | ResNet-18 (W4A4) | 6000 | 43 | | RTX4080-Laptop (12G) | 1024 | CIFAR-10 | ResNet-18 (W4A4) | 230 | 44 | | RTX4080-Laptop (12G) | 256 | ImageNet | ResNet-18 (W6A6) | 6645 | 45 | | A6000 (48G) | 512 | ImageNet | ResNet-18 (W6A6) | 4049 | 46 | | A6000 (48G) | 1024 | ImageNet | ViT-B-32 (W8A8) | 19841 | 47 | 48 | ## 4. ASiM 49 | 50 | ### 4.1. Abstract 51 | 52 | Paper Link 53 | 54 | SRAM-based Analog Compute-in-Memory (ACiM) demonstrates promising energy efficiency for deep neural network (DNN) processing. Nevertheless, efforts to optimize efficiency frequently compromise accuracy, and this trade-off remains insufficiently studied due to the difficulty of performing full-system validation. Specifically, existing simulation tools rarely target SRAM-based ACiM and exhibit inconsistent accuracy predictions, highlighting the need for a standardized, SRAM CiM circuit-aware evaluation methodology. This paper presents ASiM, a simulation framework for evaluating inference accuracy in SRAM-based ACiM systems. ASiM captures critical effects in SRAM based analog compute in memory systems, such as ADC quantization, bit parallel encoding, and analog noise, which must be modeled with high fidelity due to their distinct behavior in charge domain architectures compared to other memory technologies. ASiM supports a wide range of modern DNN workloads, including CNN and Transformer-based models such as ViT, and scales to large-scale tasks like ImageNet classification. Our results indicate that bit-parallel encoding can improve energy efficiency with only modest accuracy degradation; however, even 1 LSB of analog noise can significantly impair inference performance, particularly in complex tasks such as ImageNet. To address this, we explore hybrid analog–digital execution and majority voting schemes, both of which enhance robustness without negating energy savings. ASiM bridges the gap between hardware design and inference performance, offering actionable insights for energy-efficient, high-accuracy ACiM deployment. 55 | 56 | ### 4.2. Mode Setting for ASiM Modules 57 | Users can import and replace ASiM modules with default nn.Module into their own DNN models. ASiM modules have additional arguments in module initialization, and all configurable arguments are extracted into `config.py` within this framework. The meaning of each argument can be found in this configuration file. ASiM have two main modes in general use: **Train** and **Simulation**, for model training/fine-tuning and running bit-wise simulation. 58 | 59 | #### 4.2.1. Running Simulations by ASiM 60 | 61 | For running simulation, a pre-trained model weight has to be specified in `config.py`: `cfg.pretrain_model_path = 'model_path.pkl'`. 62 | 63 | The following example shows an ASiM module configuration for simulation: 64 | 65 | ``` 66 | ASiMConv2d(in_planes=64, 67 | out_planes=128, 68 | kernel_size=3, 69 | stride=1, 70 | padding=1, 71 | bias=False, 72 | wbit=8, 73 | xbit=8, 74 | adc_prec=8, 75 | nrow=256, 76 | rand_noise_sigma=0.05, 77 | non_linear_sigma=0.05, 78 | act_enc=None, 79 | signed_act=False, 80 | hybrid_levels=None, 81 | mode='Simulation', 82 | trim_noise=0.0, 83 | device='cuda') 84 | ``` 85 | 86 | Please note that the mode must be set as `mode='Simulation'` to switch ASiM modules into simulation mode. 87 | 88 | #### 4.2.2. Training DNN Models by ASiM 89 | 90 | The following example shows an ASiM module configuration for model training, with mode set as `mode='Train'`: 91 | 92 | ``` 93 | ASiMConv2d(in_planes=64, 94 | out_planes=128, 95 | kernel_size=3, 96 | stride=1, 97 | padding=1, 98 | bias=False, 99 | wbit=8, 100 | xbit=8, 101 | adc_prec=8, 102 | nrow=256, 103 | rand_noise_sigma=0.0, 104 | non_linear_sigma=0.0, 105 | act_enc=None, 106 | signed_act=False, 107 | hybrid_levels=None, 108 | mode='Train', 109 | trim_noise=0.0, 110 | device='cuda') 111 | ``` 112 | 113 | The `wbit` and `xbit` have to be specified for quantization-aware training. For training tiny scale DNN models (ResNet-18) with simple tasks (CIFAR-10/CIFAR-100), users can directly train the model quickly. For relately larger models (ViT) and ImageNet task, we recommend to download and process PyTorch official model weight for fine-tuning, with pretrained model weight directory specified as `cfg.pretrain_model_path = 'model_path.pkl'`. 114 | 115 | #### 4.2.3. Fine-tuning a Robust DNN Model for ACiM Applications 116 | 117 | To perform noise-aware training for obtaining a more robust DNN model, users can configure ASiM modules like: 118 | 119 | ``` 120 | ASiMConv2d(in_planes=64, 121 | out_planes=128, 122 | kernel_size=3, 123 | stride=1, 124 | padding=1, 125 | bias=False, 126 | wbit=8, 127 | xbit=8, 128 | adc_prec=8, 129 | nrow=256, 130 | rand_noise_sigma=0.0, 131 | non_linear_sigma=0.0, 132 | act_enc=None, 133 | signed_act=False, 134 | hybrid_levels=None, 135 | mode='Train', 136 | trim_noise=50.0, 137 | device='cuda') 138 | ``` 139 | 140 | Here, a pretrained model weight must be given as `cfg.pretrain_model_path = 'model_path.pkl'`, and training noise intensity must be given like `trim_noise=50.0`. After the noise-aware training, the DNN model can tolerate a certain level of analog noise. 141 | 142 | ### 4.3. Download Our Post-NAT-Models 143 | 144 | | File Name | Dataset | Model | Link | 145 | |------------------------------------------|----------|------------------|------------------------------------------------------------------------------------------------------------------------------------------| 146 | | resnet18_cifar10_w8a8_pact_trim_100.pkl | CIFAR-10 | ResNet-18 (W8A8) | Google Drive | 147 | | resnet18_imagenet_w8a8_pact_trim_50.pkl | ImageNet | ResNet-18 (W8A8) | Google Drive | 148 | | vitb32_cifar10_w8a8_trim_60.pkl | CIFAR-10 | ViT-B-32 (W8A8) | Google Drive | 149 | | vitb32_imagenet_w8a8_trim_40.pkl | ImageNet | ViT-B-32 (W8A8) | Google Drive | 150 | 151 | ## 5. Guiding Setup by Reproducing Results in the Paper 152 | 153 | In this section, we will give the setup guide by reproducing some results reported in the paper. Please follow the steps and familiarize how ASiM works. 154 | 155 | ### 5.1. Prepare Environment and Datasets 156 | 157 | Please refer our software versions provided in Sec. 2 for environment settings. Install easydict if you do not have it, and download our model weights provided in Sec. 4.3. 158 | 159 | Download the following datasets: CIFAR-10, CIFAR-100, and ImageNet, respectively. Please use `./bins/valid_setup.py` to reformat the validation folder to match our dataset file of `./dataset/imagenet.py`. 160 | 161 | The recollected dataset path should align with the following form: 162 | 163 | CIFAR-10: 164 | 165 | ``` 166 | |-- CIFAR10 167 | | |-- train 168 | | |-- test 169 | ``` 170 | 171 | CIFAR-100: 172 | 173 | ``` 174 | |-- CIFAR100 175 | | |-- train_images 176 | | |-- test_images 177 | | |-- train.csv 178 | | |-- test.csv 179 | ``` 180 | 181 | ImageNet: 182 | 183 | ``` 184 | |-- ImageNet 185 | | |-- train 186 | | | |-- n01440764 187 | | | |-- ... 188 | | |-- valid 189 | | | |-- n01440764 190 | | | |-- ... 191 | ``` 192 | 193 | ### 5.2. Simulation Examples 194 | 195 | #### 5.2.1. ResNet-18 on CIFAR-10 196 | 197 | For simulations of ResNet-18 on CIFAR-10, please refer configuration at `./examples/resnet18_cifar10/config.py`. Please give the dataset path of validation set to `cfg.test_dir`, and pre-trained model weight to `cfg.pretrain_model_path`. Based on your devices, please set the appropriate batch size and work loader to `cfg.test_bs` and `cfg.test_workers`. 198 | 199 | The default configuration is based on an 8b/8b model simulation on a bit-serial ACiM macro with row-parallelsim of 255. The CONV and Linear layer are configured separately. 200 | 201 | The model quantization bit can be modified by `wbit` and `xbit` parameters. The row-parallelsim and ADC precision can be modified by `adc_prec` and `nrow` parameters. For analog noise injection, please give `rand_noise_sigma` and `non_linear_sigma` based on your Silicon measurement. For bit-parallel scheme simulation, configure `act_enc` parameter with the encoding bit width (2, 4, etc.). To verify the efficacy of HCiM, configure the `cfg.asim_cnn_hybrid_levels` parameter with hybrid boundary level (1, 2, 3, etc.) to prevent MSB cycles from noise impact. 202 | 203 | #### 5.2.2. ResNet-18 on ImageNet 204 | 205 | Similarly, please set dataset/model weight path and refer `./examples/resnet18_imagenet/config.py` for configurations on ImageNet. 206 | 207 | #### 5.2.3. ViT-B-32 on CIFAR-10 208 | 209 | In Transformer models, the configuration of MultiheadAttention, MLP, and FC layers are given separately. Notably, in MultiheadAttention layer, we further give individual settings to QK, AV, and WQKV/WO projection, respectively. Please note the Keys (K) and Activations in projection layers are have a sign bit that have to be processed individually in bit-serial manner, we set `attn_qk_kbit` and `attn_proj_xbit` to 9 so that the LSBs can be evenly divided by 2 and 4, respectively. The HCiM boundary level in ViT is in `cfg.asim_vit_hybrid_levels`. Please refer `./examples/vitb32_cifar10/config.py` for more details. 210 | 211 | #### 5.2.4. ViT-B-32 on ImageNet 212 | 213 | Similarly, please set dataset/model weight path and refer `./examples/vitb32_imagenet/config.py` for configurations on ImageNet. 214 | 215 | #### 5.2.5. Run Simulations 216 | 217 | After the above configuration settings, just run `./main/src_simulation.py` for ASiM simulations. 218 | 219 | ### 5.3. Training Examples 220 | 221 | We give an example of training ResNet-18 model on CIFAR-10 dataset for ACiM applications. 222 | 223 | #### 5.3.1. Quantization-Aware Training (QAT) 224 | 225 | For model training, the `cfg.train_dir` and `cfg.valid_dir` have to be set as your dataset root. To perform training without an pre-trained model, just give `cfg.pretrain_model_path` a **None**. If you wish to start by a pre-trained state, please give the directory of pre-trained model weight. Please note that you may have to rearrange the parameters from original PyTorch provided models, in these cases, the state_dict have to be reformed (refer scripts in `./bins`). 226 | 227 | Please set batch size and load workers by `cfg.train_bs`, `cfg.valid_bs`, and `cfg.train_workers` respectively. Also, the `wbit` and `xbit` need to be configured to quantization weight and activation bit width, respectively. 228 | 229 | #### 5.3.2. Noise-Aware Training (NAT) 230 | 231 | For NAT, please give a pre-trained model to `cfg.pretrain_model_path`, and configure the noise intensity to `trim_noise` parameter. 232 | 233 | #### 5.3.3. Train Models 234 | 235 | After the above configuration settings, just run `./main/src_train.py` for model training. 236 | 237 | Kindly remind that training conditions such as epoch/learning rate may need to be configured case by case (`cfg.max_epoch` and `cfg.lr_init`). A `run` directory will be generated when the training is finished that containing all training log and model weights. 238 | 239 | ## 6. Model Extension 240 | 241 | As ACiM technology rapidly evolves, novel designs continue to enhance area efficiency, energy efficiency, and throughput. Although ASiM is primarily developed for charge domain SRAM ACiM, we encourage users to extend the model by tailoring the code in ASiM modules to reflect their novel circuit designs, ensuring alignment with their specific hardware characteristics. Here, we give two additional examples using ASiM as baseline to show the model's flexibility. 242 | 243 | ### 6.1. RRAM Crossbar 244 | 245 | In the case of RRAM, where the crossbar cells encode multi-bit weights directly, the loop over binary weight bits in the original simulation algorithm is unnecessary and can be omitted. Additionally, RRAM circuits are primarily affected by conductance variation, making the original capacitor nonlinearity model for SRAM-based ACiM unsuitable. Instead, we model conductance drift as a random perturbation characterized by a percentage-based standard deviation. With these minor adjustments, the ASiM framework can effectively simulate RRAM-based architectures with satisfactory accuracy. The implementation example of RRAM crossbar ACIM can be found in `./extensions/rram`. 246 | 247 | ### 6.2. Full Weight Bit-Parallel 248 | 249 | By adjusting the order of ideal MAC computing and noise injection & ADC quantization, the ASiM can be extended to further support weight bit-parallel scheme, which is a more aggressive strategy to boost energy efficiency. The implementation example of an full bit-parallel 4b/4b CNN can be found in `./extensions/w4a4_weight_full_bit_parallel`. 250 | 251 | ### 6.3. Current Domain SRAM ACiM 252 | 253 | For current domain SRAM ACiM, replacing capacitance mismatch based non-linearity model with gain compression can effectively capture the non-linearity introduced by current roll-off when transistors transition from saturation to the triode region. 254 | 255 | MAC_REAL = MAC_IDEAL / (1 + Alpha * MAC_IDEAL) 256 | 257 | Here, MAC_IDEAL represents the ideal MAC value, while MAC_REAL denotes the actual computed value. In current domain SRAM ACiM, the discharge current of the bit-line or sampling capacitor depends on the number of active transistors where the local dot product is ‘1’. When a large number of cell transistors simultaneously discharge the bit-line or capacitor, the rapid drop in Vds shifts the operation into the triode region, leading to current roll-off and introducing non-linearity into MAC computation. To replicate this behavior in simulations, we introduce an empirical parameter Alpha to align the non-linearity. By calibrating Alpha, the real MAC value remains nearly unchanged for small MAC values, while for larger MAC values, the real MAC output decreases due to gain compression effects. Additionally, incorporating random variations into binary weight tensors allows ASiM to accurately model transistor variation in current domain computations. The implementation example of current domain SRAM ACIM in CNN can be found in `./extensions/current_domain`. 258 | 259 | ## Citation 260 | 261 | If you find this repo is useful, please cite our paper. Thanks. 262 | 263 | ```bibtex 264 | @article{zhang2025asim, 265 | title={ASiM: Modeling and Analyzing Inference Accuracy of SRAM-Based Analog CiM Circuits}, 266 | author={Zhang, Wenlun and Ando, Shimpei and Chen, Yung-Chin and Yoshioka, Kentaro}, 267 | journal={IEEE Transactions on Very Large Scale Integration (VLSI) Systems}, 268 | number={01}, 269 | pages={1--14}, 270 | year={2025}, 271 | publisher={IEEE Computer Society} 272 | } 273 | ``` 274 | --------------------------------------------------------------------------------