def test_lite_module_wraps(): """Test that the wrapped module is accessible via the property.""" module = Mock() assert _LiteModule(module, Mock()).module is module wrapped_module = Mock() original_module = Mock() assert _LiteModule( wrapped_module, Mock(), original_module=original_module).module is original_module
def test_lite_module_attribute_lookup(): """Test that attribute lookup passes through to the original model when possible.""" class OriginalModule(torch.nn.Module): def __init__(self): super().__init__() self.layer = torch.nn.Linear(2, 3) self.attribute = 1 def method(self): return 2 original_module = OriginalModule() class ModuleWrapper(torch.nn.Module): def __init__(self): super().__init__() self.wrapped = original_module wrapped_module = ModuleWrapper() lite_module = _LiteModule(wrapped_module, Mock(), original_module=original_module) assert lite_module.attribute == 1 assert lite_module.layer is original_module.layer assert lite_module.method() == 2 assert lite_module.forward.__self__.__class__ == _LiteModule with pytest.raises(AttributeError): _ = lite_module.not_exists
def setup( self, model: nn.Module, *optimizers: Optimizer, move_to_device: bool = True, ) -> Any: # no specific return because the way we want our API to look does not play well with mypy """Setup a model and its optimizers for accelerated training. Args: model: A model to setup *optimizers: The optimizer(s) to setup (no optimizers is also possible) move_to_device: If set ``True`` (default), moves the model to the correct device. Set this to ``False`` and alternatively use :meth:`to_device` manually. Returns: The tuple of the wrapped model and list of optimizers, in the same order they were passed in. """ self._validate_setup(model, optimizers) if move_to_device: model = self._move_model_to_device(model=model, optimizers=list(optimizers)) # Let accelerator/plugin wrap and connect the models and optimizers model, optimizers = self._strategy._setup_model_and_optimizers(model, list(optimizers)) model = _LiteModule(model, self._precision_plugin) optimizers = [_LiteOptimizer(optimizer=optimizer, strategy=self._strategy) for optimizer in optimizers] self._models_setup += 1 if optimizers: # join both types in a list for API convenience return [model] + optimizers # type: ignore return model
def test_lite_module_forward_conversion(precision, input_type, expected_type): """Test that the LiteModule performs autocasting on the input tensors and during forward().""" lite = EmptyLite(precision=precision, accelerator="gpu", devices=1) device = torch.device("cuda", 0) def check_autocast(forward_input): assert precision != 16 or torch.is_autocast_enabled() return forward_input module = Mock(wraps=torch.nn.Identity(), side_effect=check_autocast) lite_module = _LiteModule(module, lite._precision_plugin).to(device) out = lite_module(torch.tensor([1, 2, 3], dtype=input_type, device=device)) assert module.call_args[0][0].dtype == expected_type assert out.dtype == input_type or out.dtype == torch.get_default_dtype()
def test_lite_module_device_dtype_propagation(device, dtype): """Test that the LiteModule propagates device and dtype properties to its submodules (e.g. torchmetrics).""" class DeviceModule(DeviceDtypeModuleMixin): pass device_module = DeviceModule() lite_module = _LiteModule(device_module, Mock()) lite_module.to(device) assert device_module.device == device assert lite_module.device == device lite_module.to(dtype) assert device_module.dtype == dtype assert lite_module.dtype == dtype