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