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(self):
"""Test moving AnalogSequential to cuda (from CPU)."""
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.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_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 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)
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)
