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 test_sequential_move_to_cuda_via_to(self): """Test moving AnalogSequential to cuda (from CPU), using ``.to()``.""" if not cuda.is_compiled(): raise SkipTest('not compiled with CUDA support') # Map the original tile classes to the expected ones after `cuda()`. tile_classes = { tiles.AnalogTile: tiles.CudaAnalogTile, tiles.CudaAnalogTile: tiles.CudaAnalogTile } layer = self.get_layer() expected_class = tile_classes[layer.analog_tile.tile.__class__] expected_device = device('cuda', current_device()) # Create a container and move to cuda. model = AnalogSequential(layer) model.to(device('cuda')) analog_tile = layer.analog_tile self.assertEqual(analog_tile.device, expected_device) self.assertEqual(analog_tile.get_analog_ctx().data.device, expected_device) if analog_tile.shared_weights is not None: self.assertEqual(analog_tile.shared_weights.data.device, expected_device) self.assertEqual(analog_tile.shared_weights.data.size()[0], analog_tile.tile.get_x_size()) self.assertEqual(analog_tile.shared_weights.data.size()[1], analog_tile.tile.get_d_size()) # Assert the tile has been moved to cuda. self.assertIsInstance(layer.analog_tile.tile, expected_class)
def test_sequential_move_to_cpu_via_to(self): """Test moving AnalogSequential to CPU (from CPU), using ``.to()``.""" layer = self.get_layer() # Create a container and move to cuda. model = AnalogSequential(layer) model.to(device('cpu')) # Assert the tile is still on CPU. self.assertIsInstance(layer.analog_tile, AnalogTile)
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(): """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 test_sequential_move_to_cuda_multiple_gpus(self): """Test moving AnalogSequential to cuda (from CPU), using ``.to()``.""" if not cuda.is_compiled(): raise SkipTest('not compiled with CUDA support') if device_count() < 2: raise SkipTest('Need at least two devices for this test') # Map the original tile classes to the expected ones after `cuda()`. tile_classes = { tiles.AnalogTile: tiles.CudaAnalogTile, tiles.CudaAnalogTile: tiles.CudaAnalogTile } # Test whether it can move to GPU with index 1 expected_device_num = 1 layer = self.get_layer() if isinstance(layer.analog_tile.tile.__class__, (tiles.CudaAnalogTile, tiles.CudaFloatingPointTile)): raise SkipTest('Layer is already on CUDA') expected_class = tile_classes[layer.analog_tile.tile.__class__] expected_device = device('cuda', expected_device_num) # Create a container and move to cuda. model = AnalogSequential(layer) model.cuda(device('cuda', expected_device_num)) analog_tile = layer.analog_tile self.assertEqual(analog_tile.device, expected_device) self.assertEqual(analog_tile.get_analog_ctx().data.device, expected_device) if analog_tile.shared_weights is not None: self.assertEqual(analog_tile.shared_weights.data.device, expected_device) self.assertEqual(analog_tile.shared_weights.data.size()[0], analog_tile.tile.get_x_size()) self.assertEqual(analog_tile.shared_weights.data.size()[1], analog_tile.tile.get_d_size()) # Assert the tile has been moved to cuda. self.assertIsInstance(layer.analog_tile.tile, expected_class)
def test_sequential_move_to_cuda(self): """Test sequential cuda.""" if not cuda.is_compiled(): raise SkipTest('not compiled with CUDA support') # Map the original tile classes to the expected ones after `cuda()`. tile_classes = { AnalogTile: CudaAnalogTile, CudaAnalogTile: CudaAnalogTile } layer = self.get_layer() expected_class = tile_classes[layer.analog_tile.__class__] # Create a container and move to cuda. model = AnalogSequential(layer) model.cuda() # Assert the tile has been moved to cuda. self.assertIsInstance(layer.analog_tile, expected_class)
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 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 test_sequential_move_to_cpu(self): """Test moving AnalogSequential to CPU (from CPU).""" layer = self.get_layer() # Create a container and move to cuda. model = AnalogSequential(layer) model.cpu() analog_tile = layer.analog_tile self.assertEqual(analog_tile.device, device('cpu')) self.assertEqual(analog_tile.get_analog_ctx().data.device, device('cpu')) if analog_tile.shared_weights is not None: self.assertEqual(analog_tile.shared_weights.data.device, device('cpu')) self.assertEqual(analog_tile.shared_weights.data.size()[0], analog_tile.tile.get_d_size()) self.assertEqual(analog_tile.shared_weights.data.size()[1], analog_tile.tile.get_x_size()) # Assert the tile is still on CPU. self.assertIsInstance(layer.analog_tile.tile, tiles.AnalogTile)
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 _model_from_proto(model_proto: Any) -> Module: layers = [] for layer_proto in model_proto.layers: if layer_proto.WhichOneof('item') == 'layer': layer_cls = InverseMappings.layers[layer_proto.layer.id] layer = Mappings.layers[layer_cls].from_proto( layer_proto.layer, layer_cls) else: layer_cls = InverseMappings.activation_functions[ layer_proto.activation_function.id] layer = Mappings.activation_functions[layer_cls].from_proto( layer_proto.activation_function, layer_cls) layers.append(layer) return AnalogSequential(*layers)
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_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 = 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 test_save_with_cuda(self): """Whether model is correctly reconstructed after saving""" if not cuda.is_compiled(): raise SkipTest('not compiled with CUDA support') # Map the original tile classes to the expected ones after `cuda()`. tile_classes = { tiles.AnalogTile: tiles.CudaAnalogTile, tiles.CudaAnalogTile: tiles.CudaAnalogTile } layer = self.get_layer() model = AnalogSequential(layer) model.cuda() with TemporaryFile() as file: save(model.state_dict(), file) # Create a new model and load its state dict. file.seek(0) checkpoint = load(file) model.load_state_dict(checkpoint) expected_device = device('cuda', current_device()) expected_class = tile_classes[layer.analog_tile.tile.__class__] analog_tile = model[0].analog_tile self.assertEqual(analog_tile.device, expected_device) self.assertEqual(analog_tile.get_analog_ctx().data.device, expected_device) if analog_tile.shared_weights is not None: self.assertEqual(analog_tile.shared_weights.data.device, expected_device) self.assertEqual(analog_tile.shared_weights.data.size()[0], analog_tile.tile.get_x_size()) self.assertEqual(analog_tile.shared_weights.data.size()[1], analog_tile.tile.get_d_size()) # Assert the tile has been moved to cuda. self.assertIsInstance(layer.analog_tile.tile, expected_class)
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))
def test_load_state_load_rpu_config_sequential(self): """Test creating a new model using a state dict, while using a different RPU config.""" # Create the device and the array. rpu_config_org = self.get_rpu_config() # Skipped for FP if isinstance(rpu_config_org, FloatingPointRPUConfig): raise SkipTest('Not available for FP') rpu_config_org.forward.is_perfect = False old_value = 0.11 rpu_config_org.forward.inp_noise = old_value model = AnalogSequential(self.get_layer(rpu_config=rpu_config_org)) state_dict = model.state_dict() rpu_config = deepcopy(rpu_config_org) new_value = 0.51 rpu_config.forward.inp_noise = new_value # Test restore_rpu_config=False new_model = AnalogSequential(self.get_layer(rpu_config=rpu_config)) new_model.load_state_dict(state_dict, load_rpu_config=False) new_analog_tile = self.get_analog_tile(new_model[0]) parameters = new_analog_tile.tile.get_parameters() self.assertAlmostEqual(parameters.forward_io.inp_noise, new_value) # Test restore_rpu_config=True new_model = AnalogSequential(self.get_layer(rpu_config=rpu_config)) new_model.load_state_dict(state_dict, load_rpu_config=True) new_analog_tile = self.get_analog_tile(new_model[0]) parameters = new_analog_tile.tile.get_parameters() self.assertAlmostEqual(parameters.forward_io.inp_noise, old_value)
def VGG8(): """VGG8 inspired analog model.""" model = AnalogSequential( AnalogConv2d(in_channels=3, out_channels=128, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(256), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(512), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), nn.Flatten(), AnalogLinear(in_features=8192, out_features=1024, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogLinear(in_features=1024, out_features=N_CLASSES, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.LogSoftmax(dim=1)) return model