Support for layer-specific RPUConfig: Can I assign different device types to specific layers?

[xsu0960]

Hi AIHWKIT team,

I am currently using aihwkit for training and inference. I have a question regarding the device configuration: Is it possible to specify different memory devices for different layers within the same model?

For example, can I configure some layers to use PCM-based tiles while others use ReRAM-based tiles?

Thank you for your help!

[PabloCarmona]

Hello @xsu0960! Yes, you can pass different RPU configs to each layer on the model you're building, like in the following example:

# Define one RPU configuration per layer. Each config can use a different
# device model and device parameters.
constant_step_config = SingleRPUConfig(device=ConstantStepDevice())
soft_bounds_config = SingleRPUConfig(device=SoftBoundsDevice(w_min=-0.6, w_max=0.6))

model = AnalogSequential(
    AnalogLinear(4, 2, rpu_config=constant_step_config),
    Sigmoid(),
    AnalogLinear(2, 1, rpu_config=soft_bounds_config),
)

As you saw in the example, you can pass a different rpu_config where the device is a different one. With that, your network should look something like this:

AnalogLinear(
  in_features=4, out_features=2, bias=True, SingleRPUConfig
  (analog_module): AnalogTile(RPUPulsed<float>[ConstantStep](2,4))
) 

AnalogLinear(
  in_features=2, out_features=1, bias=True, SingleRPUConfig
  (analog_module): AnalogTile(RPUPulsed<float>[SoftBounds](1,2))
)

In case you want to know more about RPU Configurations and how to use them, check our docs about it here: https://aihwkit.readthedocs.io/en/latest/using_simulator.html#rpu-configurations
I hope this answers your question. If you have further doubts, please let us know!