def test_out_scaling_alpha_learning_columnwise(self):
"""Check if out scaling alpha are learning."""
rpu_config = InferenceRPUConfig(
mapping=MappingParameter(weight_scaling_omega=0.6,
learn_out_scaling_alpha=True,
weight_scaling_omega_columnwise=True))
analog_model = Sequential(
self.get_layer(in_channels=2,
out_channels=2,
kernel_size=4,
padding=2,
rpu_config=rpu_config),
self.get_layer(in_channels=2,
out_channels=3,
kernel_size=4,
padding=2,
rpu_config=rpu_config))
loss_func = mse_loss
y_b = randn(3, 3, 6, 6, 6)
x_b = randn(3, 2, 4, 4, 4)
if self.use_cuda:
y_b = y_b.cuda()
x_b = x_b.cuda()
initial_out_scaling_alpha_0 = analog_model[
0].analog_tile.get_out_scaling_alpha().clone()
initial_out_scaling_alpha_1 = analog_model[
1].analog_tile.get_out_scaling_alpha().clone()
self.train_model(analog_model, loss_func, x_b, y_b)
learned_out_scaling_alpha_0 = analog_model[
0].analog_tile.get_out_scaling_alpha().clone()
learned_out_scaling_alpha_1 = analog_model[
1].analog_tile.get_out_scaling_alpha().clone()
self.assertGreaterEqual(initial_out_scaling_alpha_0.numel(), 1)
self.assertIsNotNone(
analog_model[0].analog_tile.get_out_scaling_alpha().grad)
self.assertNotAlmostEqualTensor(initial_out_scaling_alpha_0,
learned_out_scaling_alpha_0)
self.assertGreaterEqual(initial_out_scaling_alpha_1.numel(), 1)
self.assertIsNotNone(
analog_model[1].analog_tile.get_out_scaling_alpha().grad)
self.assertNotAlmostEqualTensor(initial_out_scaling_alpha_1,
learned_out_scaling_alpha_1)
def test_out_scaling_alpha_learning_columnwise(self):
"""Check if out scaling alpha are learning when columnwise is True."""
loss_func = mse_loss
x_b = Tensor([[0.1, 0.2, 0.3, 0.4], [0.2, 0.4, 0.3, 0.1]])
y_b = Tensor([[0.3], [0.6]])
manual_seed(4321)
rpu_config = InferenceRPUConfig(
mapping=MappingParameter(weight_scaling_omega=0.6,
learn_out_scaling_alpha=True,
weight_scaling_omega_columnwise=True))
model = Sequential(self.get_layer(4, 2, rpu_config=rpu_config),
self.get_layer(2, 1, rpu_config=rpu_config))
if self.use_cuda:
x_b = x_b.cuda()
y_b = y_b.cuda()
model = model.cuda()
initial_out_scaling_alpha_0 = model[
0].analog_tile.get_out_scaling_alpha().clone()
initial_out_scaling_alpha_1 = model[
1].analog_tile.get_out_scaling_alpha().clone()
self.train_model(model, loss_func, x_b, y_b)
learned_out_scaling_alpha_0 = model[
0].analog_tile.get_out_scaling_alpha().clone()
learned_out_scaling_alpha_1 = model[
1].analog_tile.get_out_scaling_alpha().clone()
self.assertGreaterEqual(initial_out_scaling_alpha_0.numel(), 1)
self.assertIsNotNone(model[0].analog_tile.get_out_scaling_alpha().grad)
self.assertNotAlmostEqualTensor(initial_out_scaling_alpha_0,
learned_out_scaling_alpha_0)
self.assertGreaterEqual(initial_out_scaling_alpha_1.numel(), 1)
self.assertIsNotNone(model[1].analog_tile.get_out_scaling_alpha().grad)
self.assertNotAlmostEqualTensor(initial_out_scaling_alpha_1,
learned_out_scaling_alpha_1)
def test_save_load_weight_scaling_omega(self):
"""Test saving and loading a device with weight scaling omega."""
# Create the device and the array.
rpu_config = SingleRPUConfig(mapping=MappingParameter(weight_scaling_omega=0.4))
model = self.get_layer(rpu_config=rpu_config)
analog_tile = self.get_analog_tile(model)
alpha = analog_tile.get_out_scaling_alpha().detach().cpu()
self.assertNotEqual(alpha, 1.0)
# Save the model to a file.
with TemporaryFile() as file:
save(model, file)
# Load the model.
file.seek(0)
new_model = load(file)
# Assert over the new model tile parameters.
new_analog_tile = self.get_analog_tile(new_model)
alpha_new = new_analog_tile.get_out_scaling_alpha().detach().cpu()
assert_array_almost_equal(array(alpha), array(alpha_new))
def test_analog_torch_optimizer(self):
"""Check analog layers with torch SGD for inference."""
loss_func = mse_loss
x_b = Tensor([[0.1, 0.2, 0.3, 0.4], [0.2, 0.4, 0.3, 0.1]])
y_b = Tensor([[0.3], [0.6]])
manual_seed(4321)
model = Sequential(
self.get_layer(4, 3),
self.get_layer(3, 1),
)
if not isinstance(model[0].analog_tile.rpu_config, InferenceRPUConfig):
return
manual_seed(4321)
rpu_config = FloatingPointRPUConfig(mapping=MappingParameter(
digital_bias=self.digital_bias))
model2 = AnalogSequential(
AnalogLinear(4, 3, rpu_config=rpu_config, bias=self.bias),
AnalogLinear(3, 1, rpu_config=rpu_config, bias=self.bias))
if self.use_cuda:
x_b = x_b.cuda()
y_b = y_b.cuda()
model = model.cuda()
model2 = model2.cuda()
initial_loss = loss_func(model(x_b), y_b)
# train with SGD
self.train_model_torch(model, loss_func, x_b, y_b)
self.assertLess(loss_func(model(x_b), y_b), initial_loss)
# train with AnalogSGD
self.train_model(model2, loss_func, x_b, y_b)
self.assertLess(loss_func(model2(x_b), y_b), initial_loss)
final_loss = loss_func(model(x_b), y_b).detach().cpu().numpy()
final_loss2 = loss_func(model2(x_b), y_b).detach().cpu().numpy()
assert_array_almost_equal(final_loss, final_loss2)
def __init__(
self,
in_features: int,
out_features: int,
bias: bool,
realistic_read_write: bool = False,
mapping: Optional[MappingParameter] = None,
) -> None:
# pylint: disable=super-init-not-called
self._analog_tile_counter = 0
self._registered_helper_parameter = [] # type: list
self._load_rpu_config = True
if mapping is None:
mapping = MappingParameter()
self.use_bias = bias
self.digital_bias = bias and mapping.digital_bias
self.analog_bias = bias and not mapping.digital_bias
self.realistic_read_write = realistic_read_write
self.in_features = in_features
self.out_features = out_features
# Path to store datasets
PATH_DATASET = os.path.join(os.getcwd(), 'data', 'DATASET')
# Path to store results
RESULTS = os.path.join(os.getcwd(), 'results', 'RESNET')
WEIGHT_PATH = os.path.join(RESULTS, 'example_18_model_weight.pth')
# Training parameters
SEED = 1
N_EPOCHS = 30
BATCH_SIZE = 128
LEARNING_RATE = 0.1
N_CLASSES = 10
# Device used in the RPU tile
mapping = MappingParameter(weight_scaling_omega=0.6)
RPU_CONFIG = TikiTakaEcRamPreset(mapping=mapping)
class ResidualBlock(nn.Module):
"""Residual block of a residual network with option for the skip connection."""
def __init__(self, in_ch, hidden_ch, use_conv=False, stride=1):
super().__init__()
self.conv1 = nn.Conv2d(in_ch,
hidden_ch,
kernel_size=3,
padding=1,
stride=stride)
self.bn1 = nn.BatchNorm2d(hidden_ch)
self.conv2 = nn.Conv2d(hidden_ch, hidden_ch, kernel_size=3, padding=1)
PATH_DATASET = os.path.join('data', 'DATASET')
# Path to store results
RESULTS = os.path.join(os.getcwd(), 'results', 'VGG8')
# Training parameters
SEED = 1
N_EPOCHS = 20
BATCH_SIZE = 128
LEARNING_RATE = 0.1
N_CLASSES = 10
WEIGHT_SCALING_OMEGA = 0.6 # Should not be larger than max weight.
# Select the device model to use in the training. In this case we are using one of the preset,
# but it can be changed to a number of preset to explore possible different analog devices
mapping = MappingParameter(weight_scaling_omega=WEIGHT_SCALING_OMEGA)
RPU_CONFIG = GokmenVlasovPreset(mapping=mapping)
def load_images():
"""Load images for train from torchvision datasets."""
mean = Tensor([0.4377, 0.4438, 0.4728])
std = Tensor([0.1980, 0.2010, 0.1970])
print(f'Normalization data: ({mean},{std})')
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(mean, std)])
train_set = datasets.SVHN(PATH_DATASET,
download=True,
# Imports from aihwkit.
from aihwkit.nn.conversion import convert_to_analog_mapped
from aihwkit.simulator.presets import TikiTakaReRamSBPreset
from aihwkit.simulator.configs.utils import MappingParameter
# Example: Load a predefined model from pytorch library and convert to
# its analog version.
# Load a pytorch model.
model = resnet34()
print(model)
# Define device and chip configuration used in the RPU tile
mapping = MappingParameter(
max_input_size=512, # analog tile size
max_output_size=512,
digital_bias=True,
weight_scaling_omega=0.6) # whether to use analog or digital bias
# Choose any preset or RPU configuration here
rpu_config = TikiTakaReRamSBPreset(mapping=mapping)
# Convert the model to its analog version.
# this will replace ``Linear`` layers with ``AnalogLinearMapped``
model = convert_to_analog_mapped(model, rpu_config)
# Note: One can also use ``convert_to_analog`` instead to convert
# ``Linear`` to ``AnalogLinear`` (without mapping to multiple tiles)
print(model)
def get_rpu_config(self, **kwargs):
kwargs.setdefault('mapping', MappingParameter(max_input_size=4, max_output_size=3))
return super().get_rpu_config(**kwargs)