aihwkit
aihwkit copied to clipboard
For convolutional layers, `convert_to_analog` and `convert_to_analog_mapped` behave differently
Description
For convolutional layers, convert_to_analog and convert_to_analog_mapped behave differently. For linear layers, they behave the same.
This is unexpected behavior considering the following change log entry: https://github.com/IBM/aihwkit/blob/d993b5ae50ac4d61dec0854de66812a439bc987b/CHANGELOG.md?plain=1#L120-L121
MWE:
import torch
import torch.nn as nn
import torch.nn.functional as F
from aihwkit.simulator.configs import InferenceRPUConfig
from aihwkit.nn.conversion import convert_to_analog, convert_to_analog_mapped
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
if __name__ == "__main__":
model = Net()
rpu_config = InferenceRPUConfig()
rpu_config.mapping.max_input_size = 10
rpu_config.mapping.max_output_size = 10
model_1 = convert_to_analog(model, rpu_config)
print('convert_to_analog - conv1')
for tile in model_1.conv1.analog_tiles():
print("\t", tile)
print('convert_to_analog - fc2')
for tile in model_1.fc2.analog_tiles():
print("\t", tile)
model_2 = convert_to_analog_mapped(model, rpu_config)
print('convert_to_analog_mapped - conv1')
for tile in model_2.conv1.analog_tiles():
print("\t", tile)
print('convert_to_analog_mapped - fc2')
for tile in model_2.fc2.analog_tiles():
print("\t", tile)
That's a known issue. They are not supposed to be the same. The mapped version is deprecated but the conv layer is not yet adapted to the tile module array if I remember correctly.
We are currently conducting some tests and investigating to adapt the new Conv2d layer to our tile modules