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_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)