Simulate drift after in-situ training

[alexxchen]

Example 31 shows how to use model.drift_analog_weights() with InferenceRPUConfig. But I didn't find example for SingleRPUConfig (which I used for in-situ training). I tried to modify the rpu_config into

drift_params = DriftParameter(nu=0.5)
rpu_config = SingleRPUConfig(device=ReRamSBPresetDevice(drift=drift_param))

But I can't see the drift of weights along time.

Does in-situ training with ReRamSBPresetDevice support modeling drift during inference? If yes, could you provide an example?

[jubueche]

Can you give me a minimal working example that I can quickly run?

[alexxchen]

Can you give me a minimal working example that I can quickly run?

This is my code modified based on example 31_custom_drift_models

from numpy import asarray

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

from aihwkit.nn import AnalogLinear
from aihwkit.simulator.parameters import DriftParameter
from aihwkit.simulator.presets import ReRamSBPresetDevice, ReRamSBPreset

drift_params = DriftParameter(nu=0.5)
rpu_config = ReRamSBPreset(device=ReRamSBPresetDevice())

# 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)
x = ones(BATCH_SIZE, IN_FEATURES)

# set weights
for _, layer in model.named_analog_layers():
    layer.set_weights(weights.T, zeros(OUT_FEATURES))

# Move the model and tensors to cuda if it is available
if cuda.is_compiled():
    x = x.cuda()
    model = model.cuda()

# Eval model at different time steps
model.eval()
out_lst = []
for t_inference in t_inference_times:
    model.drift_analog_weights(t_inference)  # generate new nu coefficients at each time step
    out = model(x)
    out_lst.append(out)

# Plot drift compensated outputs as a function of time
t = asarray(t_inference_times)
out_mean = asarray([mean(out).detach().cpu().numpy() for out in out_lst])
out_std = asarray([std(out).detach().cpu().numpy() for out in out_lst])
plt.figure()
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.xlabel(r"$Time$")
plt.ylabel(r"$Drift \ Compensated \ Outputs \ [1]$")
plt.tight_layout()
plt.savefig('custom_drift_model_output.png')
plt.close()

I want to know how to simulate drift on AnalogTile just as on InferenceTile.

[jubueche]

Thanks. This does not seem to be supported at the moment, sorry. @maljoras, maybe you know an elegant way to resolve this?

[maljoras]

Hi @alexxchen,

Inference type drift is different from the in-situ training in aihwkit. In the in-situ training, it is indeed assumed that long-term drift does not matter, as weight updates are made on a much smaller time-scale. That said, I agree that there is a use case of trying to estimate the effect of long-term drift after analog in-situ training. While the direct conversion of a in-situ training RPUConfig to InferenceRPUConfig is not supported yet (as there are many details how the in-situ training could be implemented), what you can do is to extract the full weights (without realistic read) from the trained model and then set it in the inference model. So something like

from aihwkit.nn import AnalogWrapper
rpu_config = ReRamSBPreset(device=ReRamSBPresetDevice()
analog_model = create_my_model(rpu_config= rpu_config) # your model creation function 
# [...] train the analog_model using analog in-situ training

# after training:
inference_rpu_config = InferenceRPUConfig(mapping=rpu_config.mapping, forward = rpu_config.forward)
# see example 34 for additional settings to get ReRAM like inference noise
inference_model = create_model(rpu_config=inference_rpu_config) 

# make sure that you have the analog wrapper utility 
weight_dic = AnalogWrapper(analog_model).get_weights() 
inference_model = AnalogWrapper(inference_model).set_weights(weight_dic)

# now use the inference_model for estimating drift
inference_model.drift_analog_weights()

Note however, that with this approach you have to make sure that the architecture is defined similarly. e.g. the rpu_config.mapping and rpu_config.forward fields etc should match in both rpu_config. Otherwise the assumed hardware architecture is different.