def create_analog_network(input_size, hidden_sizes, output_size):
"""Create the neural network using analog and digital layers.
Args:
input_size (int): size of the Tensor at the input.
hidden_sizes (list): list of sizes of the hidden layers (2 layers).
output_size (int): size of the Tensor at the output.
Returns:
nn.Module: created analog model
"""
model = AnalogSequential(
AnalogLinear(input_size,
hidden_sizes[0],
True,
rpu_config=InferenceRPUConfig()), nn.Sigmoid(),
AnalogLinear(hidden_sizes[0],
hidden_sizes[1],
True,
rpu_config=InferenceRPUConfig()), nn.Sigmoid(),
AnalogLinearMapped(hidden_sizes[1],
output_size,
True,
rpu_config=InferenceRPUConfig()),
nn.LogSoftmax(dim=1))
return model
def main():
"""Create and execute an experiment."""
model = AnalogSequential(
Flatten(),
AnalogLinear(INPUT_SIZE,
HIDDEN_SIZES[0],
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
Sigmoid(),
AnalogLinear(HIDDEN_SIZES[0],
HIDDEN_SIZES[1],
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
Sigmoid(),
AnalogLinear(HIDDEN_SIZES[1],
OUTPUT_SIZE,
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
LogSoftmax(dim=1))
# Create the training Experiment.
experiment = BasicTrainingWithScheduler(dataset=FashionMNIST,
model=model,
epochs=EPOCHS,
batch_size=BATCH_SIZE)
# Create the runner and execute the experiment.
runner = LocalRunner(device=DEVICE)
results = runner.run(experiment, dataset_root=PATH_DATASET)
print(results)
def create_analog_network(input_size, hidden_sizes, output_size):
"""Create the neural network using analog and digital layers.
Args:
input_size (int): size of the Tensor at the input.
hidden_sizes (list): list of sizes of the hidden layers (2 layers).
output_size (int): size of the Tensor at the output.
"""
model = AnalogSequential(
AnalogLinear(input_size,
hidden_sizes[0],
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
nn.Sigmoid(),
AnalogLinear(hidden_sizes[0],
hidden_sizes[1],
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
nn.Sigmoid(),
AnalogLinear(hidden_sizes[1],
output_size,
True,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
nn.LogSoftmax(dim=1))
if USE_CUDA:
model.cuda()
print(model)
return model
def __init__(self):
super().__init__()
self.feature_extractor = nn.Sequential(
AnalogConv2d(in_channels=1,
out_channels=16,
kernel_size=5,
stride=1,
rpu_config=RPU_CONFIG), nn.Tanh(),
nn.MaxPool2d(kernel_size=2),
AnalogConv2d(in_channels=16,
out_channels=32,
kernel_size=5,
stride=1,
rpu_config=RPU_CONFIG), nn.Tanh(),
nn.MaxPool2d(kernel_size=2), nn.Tanh())
self.classifier = nn.Sequential(
AnalogLinear(in_features=512,
out_features=128,
rpu_config=RPU_CONFIG),
nn.Tanh(),
AnalogLinear(in_features=128,
out_features=N_CLASSES,
rpu_config=RPU_CONFIG),
)
def create_analog_network(input_size, hidden_sizes, output_size):
"""Create the neural network using analog and digital layers.
Args:
input_size (int): size of the Tensor at the input.
hidden_sizes (list): list of sizes of the hidden layers (2 layers).
output_size (int): size of the Tensor at the output.
"""
model = nn.Sequential(
AnalogLinear(input_size,
hidden_sizes[0],
True,
resistive_device=ConstantStepResistiveDevice()),
nn.Sigmoid(),
AnalogLinear(hidden_sizes[0],
hidden_sizes[1],
True,
resistive_device=ConstantStepResistiveDevice()),
nn.Sigmoid(),
AnalogLinear(hidden_sizes[1],
output_size,
True,
resistive_device=ConstantStepResistiveDevice()),
nn.LogSoftmax(dim=1))
print(model)
return model
def __init__(
self,
input_size: int,
hidden_size: int,
bias: bool,
rpu_config: Optional[RPUConfigAlias] = None,
realistic_read_write: bool = False,
):
super().__init__()
# Default to InferenceRPUConfig
if not rpu_config:
rpu_config = InferenceRPUConfig()
self.input_size = input_size
self.hidden_size = hidden_size
self.weight_ih = AnalogLinear(
input_size,
hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
self.weight_hh = AnalogLinear(
hidden_size,
hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
def get_model(self, rpu_config: Any = TikiTakaReRamSBPreset) -> Module:
return AnalogSequential(
Flatten(),
AnalogLinear(784, 256, bias=True, rpu_config=rpu_config()),
Sigmoid(),
AnalogLinear(256, 128, bias=True, rpu_config=rpu_config()),
Sigmoid(), AnalogLinear(128,
10,
bias=True,
rpu_config=rpu_config()),
LogSoftmax(dim=1))
def __init__(self):
super().__init__()
self.dropout = nn.Dropout(DROPOUT_RATIO)
self.embedding = AnalogLinear(INPUT_SIZE,
EMBED_SIZE,
rpu_config=rpu_config)
self.lstm = AnalogLSTM(EMBED_SIZE,
HIDDEN_SIZE,
num_layers=1,
dropout=DROPOUT_RATIO,
bias=True,
rpu_config=rpu_config)
self.decoder = AnalogLinear(HIDDEN_SIZE, OUTPUT_SIZE, bias=True)
def __init__(self):
super().__init__()
self.dropout = nn.Dropout(DROPOUT_RATIO)
self.embedding = AnalogLinear(INPUT_SIZE,
EMBED_SIZE,
rpu_config=rpu_config)
self.rnn = AnalogRNN(RNN_CELL,
EMBED_SIZE,
HIDDEN_SIZE,
bidir=True,
num_layers=1,
dropout=DROPOUT_RATIO,
bias=True,
rpu_config=rpu_config)
self.decoder = AnalogLinear(2 * HIDDEN_SIZE, OUTPUT_SIZE, bias=True)
def create_analog_network():
"""Returns a Vgg8 inspired analog model."""
channel_base = 48
channel = [channel_base, 2 * channel_base, 3 * channel_base]
fc_size = 8 * channel_base
model = AnalogSequential(
nn.Conv2d(in_channels=3,
out_channels=channel[0],
kernel_size=3,
stride=1,
padding=1), nn.ReLU(),
AnalogConv2d(in_channels=channel[0],
out_channels=channel[0],
kernel_size=3,
stride=1,
padding=1,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA),
nn.BatchNorm2d(channel[0]), nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
AnalogConv2d(in_channels=channel[0],
out_channels=channel[1],
kernel_size=3,
stride=1,
padding=1,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(),
AnalogConv2d(in_channels=channel[1],
out_channels=channel[1],
kernel_size=3,
stride=1,
padding=1,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA),
nn.BatchNorm2d(channel[1]), nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
AnalogConv2d(in_channels=channel[1],
out_channels=channel[2],
kernel_size=3,
stride=1,
padding=1,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(),
AnalogConv2d(in_channels=channel[2],
out_channels=channel[2],
kernel_size=3,
stride=1,
padding=1,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA),
nn.BatchNorm2d(channel[2]), nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
nn.Flatten(),
AnalogLinear(in_features=16 * channel[2],
out_features=fc_size,
rpu_config=RPU_CONFIG,
weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(),
nn.Linear(in_features=fc_size, out_features=N_CLASSES),
nn.LogSoftmax(dim=1))
return model
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)
model2 = AnalogSequential(
AnalogLinear(4,
3,
rpu_config=FloatingPointRPUConfig(),
bias=self.bias),
AnalogLinear(3,
1,
rpu_config=FloatingPointRPUConfig(),
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,
input_size: int,
hidden_size: int,
bias: bool,
rpu_config: Optional[RPUConfigAlias],
realistic_read_write: bool = False,
weight_scaling_omega: float = 0.0
):
super().__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.weight_ih = AnalogLinear(input_size, 4 * hidden_size, bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write,
weight_scaling_omega=weight_scaling_omega)
self.weight_hh = AnalogLinear(hidden_size, 4 * hidden_size, bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write,
weight_scaling_omega=weight_scaling_omega)
def get_model(self, rpu_config: Any = ReRamSBPreset) -> Module:
return AnalogSequential(
AnalogConv2d(in_channels=3,
out_channels=16,
kernel_size=5,
stride=1,
rpu_config=rpu_config()), Tanh(),
MaxPool2d(kernel_size=2),
AnalogConv2d(in_channels=16,
out_channels=32,
kernel_size=5,
stride=1,
rpu_config=rpu_config()), Tanh(),
MaxPool2d(kernel_size=2), Tanh(), Flatten(),
AnalogLinear(in_features=800,
out_features=128,
rpu_config=rpu_config()), Tanh(),
AnalogLinear(in_features=128,
out_features=10,
rpu_config=rpu_config()), LogSoftmax(dim=1))
def create_analog_network():
"""Return a LeNet5 inspired analog model."""
channel = [16, 32, 512, 128]
model = AnalogSequential(
AnalogConv2d(in_channels=1, out_channels=channel[0], kernel_size=5, stride=1,
rpu_config=RPU_CONFIG),
nn.Tanh(),
nn.MaxPool2d(kernel_size=2),
AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=5, stride=1,
rpu_config=RPU_CONFIG),
nn.Tanh(),
nn.MaxPool2d(kernel_size=2),
nn.Tanh(),
nn.Flatten(),
AnalogLinear(in_features=channel[2], out_features=channel[3], rpu_config=RPU_CONFIG),
nn.Tanh(),
AnalogLinear(in_features=channel[3], out_features=N_CLASSES, rpu_config=RPU_CONFIG),
nn.LogSoftmax(dim=1)
)
return model
def get_model(self, rpu_config: Any = TikiTakaEcRamPreset) -> Module:
return AnalogSequential(
Conv2d(in_channels=3,
out_channels=48,
kernel_size=3,
stride=1,
padding=1), ReLU(),
AnalogConv2d(in_channels=48,
out_channels=48,
kernel_size=3,
stride=1,
padding=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), BatchNorm2d(48), ReLU(),
MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
AnalogConv2d(in_channels=48,
out_channels=96,
kernel_size=3,
stride=1,
padding=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), ReLU(),
AnalogConv2d(in_channels=96,
out_channels=96,
kernel_size=3,
stride=1,
padding=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), BatchNorm2d(96), ReLU(),
MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
AnalogConv2d(in_channels=96,
out_channels=144,
kernel_size=3,
stride=1,
padding=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), ReLU(),
AnalogConv2d(in_channels=144,
out_channels=144,
kernel_size=3,
stride=1,
padding=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), BatchNorm2d(144), ReLU(),
MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1),
Flatten(),
AnalogLinear(in_features=16 * 144,
out_features=384,
rpu_config=rpu_config(),
weight_scaling_omega=0.8), ReLU(),
Linear(in_features=384, out_features=10), LogSoftmax(dim=1))
def get_model_and_x(self):
"""Trains a simple model."""
# Prepare the datasets (input and expected output).
x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])
# Define a single-layer network, using a constant step device type.
rpu_config = self.get_rpu_config()
rpu_config.forward.out_res = -1. # Turn off (output) ADC discretization.
rpu_config.forward.w_noise_type = WeightNoiseType.ADDITIVE_CONSTANT
rpu_config.forward.w_noise = 0.02
rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)
model = AnalogLinear(4, 2, bias=True, rpu_config=rpu_config)
# Move the model and tensors to cuda if it is available.
if self.use_cuda:
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)
for _ in range(100):
opt.zero_grad()
# Add the training Tensor to the model (input).
pred = model(x)
# Add the expected output Tensor.
loss = mse_loss(pred, y)
# Run training (backward propagation).
loss.backward()
opt.step()
return model, x
def __init__(self, z_dim=10, im_dim=784, hidden_dim=128):
super().__init__()
# Build the neural network.
self.gen = AnalogSequential(
get_generator_block(z_dim, hidden_dim),
get_generator_block(hidden_dim, hidden_dim * 2),
get_generator_block(hidden_dim * 2, hidden_dim * 4),
get_generator_block(hidden_dim * 4, hidden_dim * 8),
AnalogLinear(hidden_dim * 8,
im_dim,
bias=True,
rpu_config=RPU_CONFIG),
nn.Sigmoid(),
)
def get_model(self, rpu_config: Any = EcRamPreset) -> Module:
return AnalogSequential(
AnalogConv2d(in_channels=1,
out_channels=16,
kernel_size=5,
stride=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.6), Tanh(),
MaxPool2d(kernel_size=2),
AnalogConv2d(in_channels=16,
out_channels=32,
kernel_size=5,
stride=1,
rpu_config=rpu_config(),
weight_scaling_omega=0.6), Tanh(),
MaxPool2d(kernel_size=2), Tanh(), Flatten(),
AnalogLinear(in_features=512,
out_features=128,
rpu_config=rpu_config(),
weight_scaling_omega=0.6), Tanh(),
AnalogLinear(in_features=128,
out_features=10,
rpu_config=rpu_config(),
weight_scaling_omega=0.6), LogSoftmax(dim=1))
def get_generator_block(input_dim, output_dim):
"""Return a block of the generator's neural network given input and output
dimensions.
Args:
input_dim (int): the dimension of the input vector, a scalar
output_dim (int): the dimension of the output vector, a scalar
Returns:
n.Module: a generator neural network layer, with a linear transformation
followed by a batch normalization and then a relu activation
"""
return AnalogSequential(
AnalogLinear(input_dim, output_dim, bias=True, rpu_config=RPU_CONFIG),
nn.BatchNorm1d(output_dim),
nn.ReLU(inplace=True),
)
def get_layer(self, in_features=3, out_features=4, **kwargs):
kwargs.setdefault('rpu_config', self.get_rpu_config())
kwargs.setdefault('bias', self.bias)
return AnalogLinear(in_features, out_features, **kwargs).cuda()
from torch import Tensor
from torch.nn.functional import mse_loss
# Imports from aihwkit.
from aihwkit.nn import AnalogLinear
from aihwkit.optim.analog_sgd import AnalogSGD
from aihwkit.simulator.devices import ConstantStepResistiveDevice
from aihwkit.simulator.rpu_base import cuda
# Prepare the datasets (input and expected output).
x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])
# Define a single-layer network, using a constant step device type.
analog_device = ConstantStepResistiveDevice()
model = AnalogLinear(4, 2, bias=True, resistive_device=analog_device)
# Move the model and tensors to cuda if it is available.
if cuda.is_compiled():
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)
for epoch in range(100):
# Add the training Tensor to the model (input).
pred = model(x)
# Add the expected output Tensor.
from aihwkit.simulator.configs.devices import ConstantStepDevice
from aihwkit.utils.analog_info import analog_summary
# Define a single-layer network, using a constant step device type.
rpu_config = SingleRPUConfig(device=ConstantStepDevice())
channel = [16, 32, 512, 128]
model = AnalogSequential(
AnalogConv2d(in_channels=1,
out_channels=channel[0],
kernel_size=5,
stride=1,
rpu_config=rpu_config), nn.Tanh(),
nn.MaxPool2d(kernel_size=2),
AnalogConv2d(in_channels=channel[0],
out_channels=channel[1],
kernel_size=5,
stride=1,
rpu_config=rpu_config), nn.Tanh(),
nn.MaxPool2d(kernel_size=2), nn.Tanh(), nn.Flatten(),
AnalogLinear(in_features=channel[2],
out_features=channel[3],
rpu_config=rpu_config), nn.Tanh(),
AnalogLinear(in_features=channel[3],
out_features=10,
rpu_config=rpu_config), nn.LogSoftmax(dim=1))
# TODO: add mapped examples
# print(model.__class__.__name__)
analog_summary(model, (1, 1, 28, 28))
from torch.nn.functional import mse_loss
# Imports from aihwkit.
from aihwkit.nn import AnalogLinear
from aihwkit.optim import AnalogSGD
from aihwkit.simulator.configs import SingleRPUConfig
from aihwkit.simulator.configs.devices import ConstantStepDevice
from aihwkit.simulator.rpu_base import cuda
# Prepare the datasets (input and expected output).
x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])
# Define a single-layer network, using a constant step device type.
rpu_config = SingleRPUConfig(device=ConstantStepDevice())
model = AnalogLinear(4, 2, bias=True, rpu_config=rpu_config)
# Move the model and tensors to cuda if it is available.
if cuda.is_compiled():
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)
for epoch in range(100):
# Add the training Tensor to the model (input).
pred = model(x)
# Add the expected output Tensor.
rpu_config.forward.w_noise = 0.02 # Short-term w-noise.
rpu_config.clip.type = WeightClipType.FIXED_VALUE
rpu_config.clip.fixed_value = 1.0
rpu_config.modifier.pdrop = 0.03 # Drop connect.
rpu_config.modifier.type = WeightModifierType.ADD_NORMAL # Fwd/bwd weight noise.
rpu_config.modifier.std_dev = 0.1
rpu_config.modifier.rel_to_actual_wmax = True
# Inference noise model.
rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)
# drift compensation
rpu_config.drift_compensation = GlobalDriftCompensation()
model = AnalogLinear(4, 2, bias=True, rpu_config=rpu_config)
# Move the model and tensors to cuda if it is available.
if cuda.is_compiled():
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)
print(model.analog_tile.tile)
for epoch in range(100):
# Add the training Tensor to the model (input).
from torch import Tensor
from torch.nn.functional import mse_loss
from torch.nn import Sequential
# Imports from aihwkit.
from aihwkit.nn import AnalogLinear
from aihwkit.optim.analog_sgd import AnalogSGD
from aihwkit.simulator.devices import ConstantStepResistiveDevice
# Prepare the datasets (input and expected output).
x_b = Tensor([[0.1, 0.2, 0.0, 0.0], [0.2, 0.4, 0.0, 0.0]])
y_b = Tensor([[0.3], [0.6]])
# Define a multiple-layer network, using a constant step device type.
model = Sequential(
AnalogLinear(4, 2, resistive_device=ConstantStepResistiveDevice()),
AnalogLinear(2, 2, resistive_device=ConstantStepResistiveDevice()),
AnalogLinear(2, 1, resistive_device=ConstantStepResistiveDevice()))
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.5)
opt.regroup_param_groups(model)
for epoch in range(100):
# Add the training Tensor to the model (input).
pred = model(x_b)
# Add the expected output Tensor.
loss = mse_loss(pred, y_b)
# Run training (backward propagation).
loss.backward()
def test_against_fp(self):
"""Test whether FP is same as is_perfect inference tile."""
# pylint: disable-msg=too-many-locals
# Prepare the datasets (input and expected output).
x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])
# Define a single-layer network, using a constant step device type.
rpu_config = self.get_rpu_config()
rpu_config.forward.is_perfect = True
model_torch = Linear(4, 2, bias=True)
model = AnalogLinear(4, 2, bias=True, rpu_config=rpu_config)
model.set_weights(model_torch.weight, model_torch.bias)
model_fp = AnalogLinear(4,
2,
bias=True,
rpu_config=FloatingPointRPUConfig())
model_fp.set_weights(model_torch.weight, model_torch.bias)
self.assertTensorAlmostEqual(model.get_weights()[0],
model_torch.weight)
self.assertTensorAlmostEqual(model.get_weights()[0],
model_fp.get_weights()[0])
# Move the model and tensors to cuda if it is available.
if self.use_cuda:
x = x.cuda()
y = y.cuda()
model.cuda()
model_fp.cuda()
model_torch.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt_fp = AnalogSGD(model_fp.parameters(), lr=0.1)
opt_torch = SGD(model_torch.parameters(), lr=0.1)
for _ in range(100):
# inference
opt.zero_grad()
pred = model(x)
loss = mse_loss(pred, y)
loss.backward()
opt.step()
# same for fp
opt_fp.zero_grad()
pred_fp = model_fp(x)
loss_fp = mse_loss(pred_fp, y)
loss_fp.backward()
opt_fp.step()
# same for torch
opt_torch.zero_grad()
pred_torch = model_torch(x)
loss_torch = mse_loss(pred_torch, y)
loss_torch.backward()
opt_torch.step()
self.assertTensorAlmostEqual(pred_torch, pred)
self.assertTensorAlmostEqual(loss_torch, loss)
self.assertTensorAlmostEqual(model.get_weights()[0],
model_torch.weight)
self.assertTensorAlmostEqual(pred_fp, pred)
self.assertTensorAlmostEqual(loss_fp, loss)
self.assertTensorAlmostEqual(model.get_weights()[0],
model_fp.get_weights()[0])
from torch.nn.functional import mse_loss
from torch.nn import Sequential
# Imports from aihwkit.
from aihwkit.nn import AnalogLinear
from aihwkit.optim import AnalogSGD
from aihwkit.simulator.configs import SingleRPUConfig
from aihwkit.simulator.configs.devices import ConstantStepDevice
# Prepare the datasets (input and expected output).
x_b = Tensor([[0.1, 0.2, 0.0, 0.0], [0.2, 0.4, 0.0, 0.0]])
y_b = Tensor([[0.3], [0.6]])
# Define a multiple-layer network, using a constant step device type.
model = Sequential(
AnalogLinear(4, 2,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
AnalogLinear(2, 2,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())),
AnalogLinear(2, 1,
rpu_config=SingleRPUConfig(device=ConstantStepDevice())))
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.5)
opt.regroup_param_groups(model)
for epoch in range(100):
# Add the training Tensor to the model (input).
pred = model(x_b)
# Add the expected output Tensor.
loss = mse_loss(pred, y_b)
# Run training (backward propagation).
def get_layer(self, in_features=3, out_features=4, **kwargs):
kwargs.setdefault('rpu_config', self.get_rpu_config())
kwargs.setdefault('bias', self.bias)
kwargs['rpu_config'].mapping.digital_bias = self.digital_bias
return AnalogLinear(in_features, out_features, **kwargs)
class AnalogLSTMCell(AnalogSequential):
"""Analog LSTM Cell.
Args:
input_size: in_features size for W_ih matrix
hidden_size: in_features and out_features size for W_hh matrix
bias: whether to use a bias row on the analog tile or not
rpu_config: configuration for an analog resistive processing unit
realistic_read_write: whether to enable realistic read/write
for setting initial weights and read out of weights
"""
# pylint: disable=abstract-method
def __init__(
self,
input_size: int,
hidden_size: int,
bias: bool,
rpu_config: Optional[RPUConfigAlias] = None,
realistic_read_write: bool = False,
):
super().__init__()
# Default to InferenceRPUConfig
if not rpu_config:
rpu_config = InferenceRPUConfig()
self.input_size = input_size
self.hidden_size = hidden_size
self.weight_ih = AnalogLinear(
input_size,
4 * hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
self.weight_hh = AnalogLinear(
hidden_size,
4 * hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
def get_zero_state(self, batch_size: int) -> Tensor:
"""Returns a zeroed state.
Args:
batch_size: batch size of the input
Returns:
Zeroed state tensor
"""
device = self.weight_ih.get_analog_tile_devices()[0]
return LSTMState(zeros(batch_size, self.hidden_size, device=device),
zeros(batch_size, self.hidden_size, device=device))
def forward(
self, input_: Tensor,
state: Tuple[Tensor,
Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]:
# pylint: disable=arguments-differ
h_x, c_x = state
gates = self.weight_ih(input_) + self.weight_hh(h_x)
in_gate, forget_gate, cell_gate, out_gate = gates.chunk(4, 1)
in_gate = sigmoid(in_gate)
forget_gate = sigmoid(forget_gate)
cell_gate = tanh(cell_gate)
out_gate = sigmoid(out_gate)
c_y = (forget_gate * c_x) + (in_gate * cell_gate)
h_y = out_gate * tanh(c_y)
return h_y, (h_y, c_y)
class AnalogVanillaRNNCell(AnalogSequential):
"""Analog Vanilla RNN Cell.
Args:
input_size: in_features size for W_ih matrix
hidden_size: in_features and out_features size for W_hh matrix
bias: whether to use a bias row on the analog tile or not
rpu_config: configuration for an analog resistive processing unit
realistic_read_write: whether to enable realistic read/write
for setting initial weights and read out of weights
"""
# pylint: disable=abstract-method
def __init__(
self,
input_size: int,
hidden_size: int,
bias: bool,
rpu_config: Optional[RPUConfigAlias] = None,
realistic_read_write: bool = False,
):
super().__init__()
# Default to InferenceRPUConfig
if not rpu_config:
rpu_config = InferenceRPUConfig()
self.input_size = input_size
self.hidden_size = hidden_size
self.weight_ih = AnalogLinear(
input_size,
hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
self.weight_hh = AnalogLinear(
hidden_size,
hidden_size,
bias=bias,
rpu_config=rpu_config,
realistic_read_write=realistic_read_write)
def get_zero_state(self, batch_size: int) -> Tensor:
"""Returns a zeroed state.
Args:
batch_size: batch size of the input
Returns:
Zeroed state tensor
"""
device = self.weight_ih.get_analog_tile_devices()[0]
return zeros(batch_size, self.hidden_size, device=device)
def forward(self, input_: Tensor, state: Tensor) -> Tuple[Tensor, Tensor]:
# pylint: disable=arguments-differ
igates = self.weight_ih(input_)
hgates = self.weight_hh(state)
out = tanh(igates + hgates)
return out, out # output will also be hidden state
