IBM / aihwkit

IBM Analog Hardware Acceleration Kit
https://aihwkit.readthedocs.io
Apache License 2.0
343 stars 143 forks source link

For convolutional layers, `convert_to_analog` and `convert_to_analog_mapped` behave differently #627

Open coreylammie opened 5 months ago

coreylammie commented 5 months ago

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)    
maljoras commented 5 months ago

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.

PabloCarmona commented 1 month ago

We are currently conducting some tests and investigating to adapt the new Conv2d layer to our tile modules