Skip to content

Add quantization library #719

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
Open

Add quantization library #719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
bbf4186
Added quantized library and infrastructure
Feb 17, 2025
3ccb056
Merge branch 'master' into quantized_library
Feb 17, 2025
55b5d30
Make pylint and black happy :)
atvasilopoulos Feb 17, 2025
aeb2752
Merge remote-tracking branch 'upstream/master' into quantized_library
atvasilopoulos Mar 5, 2025
9a08d06
Added example notebook
atvasilopoulos Mar 31, 2025
9fbabd7
Merge remote-tracking branch 'upstream/master' into quantized_library
atvasilopoulos Mar 31, 2025
907f43b
Make mypy/black/pylint happier
atvasilopoulos Apr 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 7 additions & 27 deletions examples/30_external_hardware_aware_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,7 @@
from torchvision.datasets.utils import download_url
from aihwkit.nn.conversion import convert_to_analog
from aihwkit.simulator.presets import StandardHWATrainingPreset
from aihwkit.inference.calibration import (
calibrate_input_ranges,
InputRangeCalibrationType,
)
from aihwkit.inference.calibration import calibrate_input_ranges, InputRangeCalibrationType


class LambdaLayer(torch.nn.Module):
Expand All @@ -48,9 +45,7 @@ def __init__(self, in_planes, planes, stride=1, option="A"):
in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(
planes, planes, kernel_size=3, stride=1, padding=1, bias=False
)
self.conv2 = torch.nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = torch.nn.BatchNorm2d(planes)

self.shortcut = torch.nn.Sequential()
Expand All @@ -61,20 +56,13 @@ def __init__(self, in_planes, planes, stride=1, option="A"):
"""
self.shortcut = LambdaLayer(
lambda x: F.pad(
x[:, :, ::2, ::2],
(0, 0, 0, 0, planes // 4, planes // 4),
"constant",
0,
x[:, :, ::2, ::2], (0, 0, 0, 0, planes // 4, planes // 4), "constant", 0
)
)
elif option == "B":
self.shortcut = torch.nn.Sequential(
torch.nn.Conv2d(
in_planes,
self.expansion * planes,
kernel_size=1,
stride=stride,
bias=False,
in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False
),
torch.nn.BatchNorm2d(self.expansion * planes),
)
Expand Down Expand Up @@ -175,20 +163,14 @@ def get_test_loader(batch_size=128):
transform_test = torchvision.transforms.Compose(
[
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
),
torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
]
)
testset = torchvision.datasets.CIFAR10(
root="data/cifar10", train=False, download=True, transform=transform_test
)
test_loader = torch.utils.data.DataLoader(
testset,
batch_size=batch_size,
shuffle=False,
num_workers=0,
pin_memory=True,
testset, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True
)
return test_loader

Expand Down Expand Up @@ -257,9 +239,7 @@ def evaluate_model(model, test_loader, device):
for t_id, t in enumerate(t_inferences):
for i in range(n_reps):
model.drift_analog_weights(t)
inference_accuracy_values[t_id, i] = evaluate_model(
model, test_loader, device
)
inference_accuracy_values[t_id, i] = evaluate_model(model, test_loader, device)

print(
f"Test set accuracy (%) at t={t}s: mean: {inference_accuracy_values[t_id].mean()}, \
Expand Down
66 changes: 31 additions & 35 deletions examples/31_custom_drift_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,7 @@
from numpy import asarray

# Imports from PyTorch.
from torch import (
zeros,
ones,
mean,
std,
linspace,
)
from torch import zeros, ones, mean, std, linspace
import matplotlib.pyplot as plt

# Imports from aihwkit.
Expand All @@ -35,19 +29,20 @@
from aihwkit.simulator.parameters.io import IOParameters
from aihwkit.simulator.configs import TorchInferenceRPUConfig

g_min, g_max = 0.0, 25.
g_min, g_max = 0.0, 25.0
# define custom drift model
custom_drift_model = dict(g_lst=[g_min, 10., g_max],
nu_mean_lst=[0.08, 0.05, 0.03],
nu_std_lst=[0.03, 0.02, 0.01])

t_inference_times = [1, # 1 sec
60, # 1 min
60 * 60, # 1 hour
24 * 60 * 60, # 1 day
30 * 24 * 60 * 60, # 1 month
12 * 30 * 24 * 60 * 60, # 1 year
]
custom_drift_model = dict(
g_lst=[g_min, 10.0, g_max], nu_mean_lst=[0.08, 0.05, 0.03], nu_std_lst=[0.03, 0.02, 0.01]
)

t_inference_times = [
1, # 1 sec
60, # 1 min
60 * 60, # 1 hour
24 * 60 * 60, # 1 day
30 * 24 * 60 * 60, # 1 month
12 * 30 * 24 * 60 * 60, # 1 year
]

IN_FEATURES = 512
OUT_FEATURES = 512
Expand All @@ -57,26 +52,27 @@
io_params = IOParameters(
bound_management=BoundManagementType.NONE,
nm_thres=1.0,
inp_res=2 ** 8 - 2,
inp_res=2**8 - 2,
out_bound=-1,
out_res=-1,
out_noise=0.0)
out_noise=0.0,
)

noise_model = CustomDriftPCMLikeNoiseModel(custom_drift_model,
prog_noise_scale=0.0, # turn off to show drift only
read_noise_scale=0.0, # turn off to show drift only
drift_scale=1.0,
g_converter=SinglePairConductanceConverter(g_min=g_min,
g_max=g_max),
)
noise_model = CustomDriftPCMLikeNoiseModel(
custom_drift_model,
prog_noise_scale=0.0, # turn off to show drift only
read_noise_scale=0.0, # turn off to show drift only
drift_scale=1.0,
g_converter=SinglePairConductanceConverter(g_min=g_min, g_max=g_max),
)

rpu_config = TorchInferenceRPUConfig(noise_model=noise_model, forward=io_params)

# define simple model, weights, and activations
model = AnalogLinear(IN_FEATURES, OUT_FEATURES, bias=False, rpu_config=rpu_config)
weights = linspace(custom_drift_model['g_lst'][0],
custom_drift_model['g_lst'][-1],
OUT_FEATURES).repeat(IN_FEATURES, 1)
weights = linspace(
custom_drift_model["g_lst"][0], custom_drift_model["g_lst"][-1], OUT_FEATURES
).repeat(IN_FEATURES, 1)
x = ones(BATCH_SIZE, IN_FEATURES)

# set weights
Expand Down Expand Up @@ -107,14 +103,14 @@
plt.xlabel(r"$Conductance \ [\mu S]$")
plt.ylabel(r"$\nu \ [1]$")
plt.tight_layout()
plt.savefig('custom_drift_model.png')
plt.savefig("custom_drift_model.png")
plt.close()

# create simple linear layer model
model = AnalogLinear(IN_FEATURES, OUT_FEATURES, bias=False, rpu_config=rpu_config)

# define weights, activations
weights = (1. / 512.) * ones(IN_FEATURES, OUT_FEATURES)
weights = (1.0 / 512.0) * ones(IN_FEATURES, OUT_FEATURES)
x = ones(BATCH_SIZE, IN_FEATURES)

# set weights
Expand Down Expand Up @@ -142,9 +138,9 @@
plt.plot(t, out_mean)
plt.fill_between(t, out_mean - out_std, out_mean + out_std, alpha=0.2)
plt.xscale("log")
plt.xticks(t, ['1 sec', '1 min', '1 hour', '1 day', '1 month', '1 year'], rotation='vertical')
plt.xticks(t, ["1 sec", "1 min", "1 hour", "1 day", "1 month", "1 year"], rotation="vertical")
plt.xlabel(r"$Time$")
plt.ylabel(r"$Drift \ Compensated \ Outputs \ [1]$")
plt.tight_layout()
plt.savefig('custom_drift_model_output.png')
plt.savefig("custom_drift_model_output.png")
plt.close()
106 changes: 57 additions & 49 deletions examples/32_weight_programming_options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@
SinglePairConductanceConverter,
DualPairConductanceConverter,
NPairConductanceConverter,
CustomPairConductanceConverter
CustomPairConductanceConverter,
)
from aihwkit.inference import PCMLikeNoiseModel
from aihwkit.simulator.configs import InferenceRPUConfig
Expand All @@ -64,22 +64,24 @@
SEGMENTS = 32
USE_CUDA = True

time_dict = {'1 second': 1,
'1 month' : 1 * 60 * 60 * 24 * 30}
time_dict = {"1 second": 1, "1 month": 1 * 60 * 60 * 24 * 30}

g_min, g_max = 0.1, 15.
g_min, g_max = 0.1, 15.0
single_pair_g_converter = SinglePairConductanceConverter(g_min=g_min, g_max=g_max)
dual_pair_g_converter = DualPairConductanceConverter(f_lst=[1.0, 3.0], g_min=g_min, g_max=g_max)
npair_g_converter = NPairConductanceConverter(f_lst=[1.0, 2.0, 3.0], g_min=g_min, g_max=g_max)

# custom programming model A (curved)
k_lst = [0.5, 0.8]
g_plus_ = [g_min] + [k * (g_max - g_min) + g_min for k in k_lst] + [g_max]
g_minus_ = [g_min] + [(k - (i / (len(k_lst) + 1))) * (g_max - g_min) + g_min
for i, k in enumerate(k_lst, 1)] + [g_min]
g_minus_ = (
[g_min]
+ [(k - (i / (len(k_lst) + 1))) * (g_max - g_min) + g_min for i, k in enumerate(k_lst, 1)]
+ [g_min]
)
g_plus = g_minus_[::-1][:-1] + g_plus_
g_minus = g_plus_[::-1][:-1] + g_minus_
prog_model = {'A': [g_plus, g_minus]}
prog_model = {"A": [g_plus, g_minus]}

# custom programming model B (random)
n_pts = 5
Expand All @@ -89,57 +91,63 @@
g_minus_ = [g - lb + g_min for g, lb in zip(g_plus_, lower_bounds)]
g_plus = g_minus_[::-1][:-1] + g_plus_
g_minus = g_plus_[::-1][:-1] + g_minus_
prog_model.update({'B': [g_plus, g_minus]})
prog_model.update({"B": [g_plus, g_minus]})

custom_g_converter = CustomPairConductanceConverter(f_lst=[1.0],
g_lst=prog_model['A'],
g_min=g_min,
g_max=g_max,
invertibility_test=True)
custom_g_converter = CustomPairConductanceConverter(
f_lst=[1.0], g_lst=prog_model["A"], g_min=g_min, g_max=g_max, invertibility_test=True
)


def plot_weights(g_converter: Type[BaseConductanceConverter],
ideal_weights: Tensor,
drifted_weights: Tensor,
suffix: str = '') -> None:
"""Plots weight programming strategy
"""
def plot_weights(
g_converter: Type[BaseConductanceConverter],
ideal_weights: Tensor,
drifted_weights: Tensor,
suffix: str = "",
) -> None:
"""Plots weight programming strategy"""

g_lst, params = g_converter.convert_to_conductances(drifted_weights)
return_weights = g_converter.convert_back_to_weights(g_lst, params)

assert allclose(drifted_weights, return_weights, atol=0.0001), \
"conversion error: weights don't match for %s" % str(g_converter)
assert allclose(
drifted_weights, return_weights, atol=0.0001
), "conversion error: weights don't match for %s" % str(g_converter)

g_converter_name = str(g_converter).split('(', maxsplit=1)[0]
g_converter_name = str(g_converter).split("(", maxsplit=1)[0]
rows = int(len(g_lst) / 2) + 1
width, height = 7, 4
plt.subplots(rows, 1, figsize=(width, rows * height))
plt.subplot(rows, 1, 1)
if suffix != 'ideal':
if suffix != "ideal":
title_str = "%s @ %s" % (g_converter_name, suffix)
else:
title_str = "Ideal %s" % g_converter_name
plt.title(title_str)
plt.plot(ideal_weights.detach().cpu().numpy().flatten(),
return_weights.detach().cpu().numpy().flatten(),
'.',
ms=1)
plt.plot(
ideal_weights.detach().cpu().numpy().flatten(),
return_weights.detach().cpu().numpy().flatten(),
".",
ms=1,
)
plt.xlabel(r"$Ideal \ Weights \ [1]$")
plt.ylabel(r"$Actual \ Weights \ [1]$")

for i, (gp, gm) in enumerate(zip(g_lst[::2], g_lst[1::2])):
plt.subplot(rows, 1, i + 2)
plt.plot(ideal_weights.detach().cpu().numpy().flatten(),
gp.detach().cpu().numpy().flatten(),
'.',
ms=1,
label=r"$G^+_%d$" % i)
plt.plot(ideal_weights.detach().cpu().numpy().flatten(),
gm.detach().cpu().numpy().flatten(),
'.',
ms=1,
label=r"$G^-_%d$" % i)
plt.plot(
ideal_weights.detach().cpu().numpy().flatten(),
gp.detach().cpu().numpy().flatten(),
".",
ms=1,
label=r"$G^+_%d$" % i,
)
plt.plot(
ideal_weights.detach().cpu().numpy().flatten(),
gm.detach().cpu().numpy().flatten(),
".",
ms=1,
label=r"$G^-_%d$" % i,
)
plt.legend()
plt.xlabel(r"$Ideal \ Weights \ [1]$")
plt.ylabel(r"$Conductance \ [\mu S]$")
Expand All @@ -149,28 +157,28 @@ def plot_weights(g_converter: Type[BaseConductanceConverter],


def main():
"""Compare weight programming strategies (i.e. g_converters)
"""
"""Compare weight programming strategies (i.e. g_converters)"""

# create dataset
x = clamp((1.0 / 3.0) * randn(BATCH_SIZE, IN_FEATURES), min=-1, max=1)
ideal_weights = clamp((1.0 / 3.0) * randn(IN_FEATURES, OUT_FEATURES), min=-1, max=1)

# compare each g_converter
for g_converter in [single_pair_g_converter,
dual_pair_g_converter,
npair_g_converter,
custom_g_converter]:
for g_converter in [
single_pair_g_converter,
dual_pair_g_converter,
npair_g_converter,
custom_g_converter,
]:

# g_converter applied to ideal weights
plot_weights(g_converter, ideal_weights, ideal_weights, 'ideal')
plot_weights(g_converter, ideal_weights, ideal_weights, "ideal")

# create simple model using g_converter
rpu_config = InferenceRPUConfig()
rpu_config.noise_model = PCMLikeNoiseModel(prog_noise_scale=1.0,
read_noise_scale=1.0,
drift_scale=1.0,
g_converter=g_converter)
rpu_config.noise_model = PCMLikeNoiseModel(
prog_noise_scale=1.0, read_noise_scale=1.0, drift_scale=1.0, g_converter=g_converter
)
model = AnalogLinear(IN_FEATURES, OUT_FEATURES, bias=False, rpu_config=rpu_config)

# set weights
Expand All @@ -182,7 +190,7 @@ def main():
for t_name, t_inference in time_dict.items():

model.drift_analog_weights(t_inference)
_ = model(x) # dummy inference applies programming errors and read noise
_ = model(x) # dummy inference applies programming errors and read noise
drifted_weights, _ = model.get_weights()

# g_converter applied to non-ideal weights
Expand Down
Loading