Repository: pytorch/ignite

Issue:
model is moved too early in create_supervised_trainer/evaluator
Hi, first of all, thanks for the great library!
I have a general "bug" report or question on the specific implementation of `create_supervised_trainer` and `create_supervised_evaluator`.
As seen for example here:
https://github.com/pytorch/ignite/blob/master/ignite/engine/__init__.py#L64

The moment a supervised evaluator/trainer is created the model is moved to the specific device and 
*not* during the inference stage.

I agree, that this version doesn't have an impact in general, but I would argue that the model should be moved in the `_inference` function and not when I create the evaluator/trainer.
A weird example but similar to my case:
I have multiple evaluators and I would like to run one on the CPU and one on the GPU.
The reason being that a different library doesn't support GPUs yet and I would like to add it in an "ignite" style to keep the code similar.
Now, the run steps will not work, as the model is *not* moved again in the inference stage.

I can imagine that there is a reason why it is the way it is, but I am wondering if a future version couldn't do the `model.to(device)` step in the `_inference` function. :)

Thanks!



Discussion:
@kai-tub thanks for awesome feedback and thanks for your remark ! I understand that this could be done in a more lazy way, for example, when engine is started (probably, calling `model.to` every iteration I think can have some performance impact.) 
Definitely, `create_supervised_*` can setup model's device in a handler:
```python
def create_supervised_trainer(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    loss_fn: Union[Callable, torch.nn.Module],
    device: Optional[Union[str, torch.device]] = None,
    non_blocking: bool = False,
    prepare_batch: Callable = _prepare_batch,
    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
) -> Engine:

    def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
        model.train()
        optimizer.zero_grad()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        loss.backward()
        optimizer.step()
        return output_transform(x, y, y_pred, loss)

    trainer = Engine(_update)

    if device:
         @trainer.on(Events.STARTED)
         def setup_device(_):
              model.to(device)

    return trainer
```

In this way, certainly we can alternate for example, training on GPU and validating on CPU etc.

@justusschock @ykumards any thoughts ? 


This looks great!
I don't see why this shouldn't be adapted (or are there different parts of the library that work similarly?)

And thanks for the quick reply. :)
@kai-tub main reason for that I would say that there were no particular need to switch devices like that. For example, if we use DDP wrapper for model, we train and evaluate within the same context. And another point is that `create_supervised*` are just basic helpers, personnally I setup by myself the update function etc.
I would rather move the model more often. This way the user is more flexible and also these calls become no-ops in PyTorch, when the model is already on the correct device
@vfdev-5 I agree that setting up the update function is the obvious alternative (and is what I've done).
But for someone who never used PyTorch half a year ago, I really liked to use ignite to learn Pytorch in a "top-down" fashion. I could focus more on the design of the workflow and less on the actual implementation. Afterwards I looked at the underlying implementation to understand how it works.  Plus I think that the best code is the code that was never written. 

I think that most cases are solved by `create_supervised*` and would be more applicable when the moving process is moved after the function initialization. :) 
Since this is labeled as "help wanted" should I issue a PR with the changes mentioned by @justusschock?
yes, please

On Tue, Mar 10, 2020, 09:47 kai-tub <notifications@github.com> wrote:

> Since this is labeled as "help wanted" should I issue a PR with the
> changes mentioned by @justusschock <https://github.com/justusschock>?
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pytorch/ignite/issues/830?email_source=notifications&email_token=AASYOHYP6NPNVYLKAVAO7BDRGX5BVA5CNFSM4LCQXVF2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEOKQP7Q#issuecomment-596969470>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AASYOH4PIKK5KUO7IVOBZ7LRGX5BVANCNFSM4LCQXVFQ>
> .
>


Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -7,6 +7,7 @@
 from ignite.metrics import Metric
 
 __all__ = [
+    "State",
     "create_supervised_trainer",
     "create_supervised_evaluator",
     "Engine",
@@ -46,7 +47,7 @@ def create_supervised_trainer(
         optimizer (`torch.optim.Optimizer`): the optimizer to use.
         loss_fn (torch.nn loss function): the loss function to use.
         device (str, optional): device type specification (default: None).
-            Applies to both model and batches.
+            Applies to batches and the model after starting the engine.
         non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
@@ -60,8 +61,6 @@ def create_supervised_trainer(
     Returns:
         Engine: a trainer engine with supervised update function.
     """
-    if device:
-        model.to(device)
 
     def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
         model.train()
@@ -73,7 +72,15 @@ def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[t
         optimizer.step()
         return output_transform(x, y, y_pred, loss)
 
-    return Engine(_update)
+    trainer = Engine(_update)
+
+    if device is not None:
+
+        @trainer.on(Events.STARTED)
+        def move_model(engine):
+            model.to(device)
+
+    return trainer
 
 
 def create_supervised_evaluator(
@@ -82,7 +89,7 @@ def create_supervised_evaluator(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
-    output_transform: Callable = lambda x, y, y_pred: (y_pred, y,),
+    output_transform: Callable = lambda x, y, y_pred: (y_pred, y),
 ) -> Engine:
     """
     Factory function for creating an evaluator for supervised models.
@@ -91,7 +98,7 @@ def create_supervised_evaluator(
         model (`torch.nn.Module`): the model to train.
         metrics (dict of str - :class:`~ignite.metrics.Metric`): a map of metric names to Metrics.
         device (str, optional): device type specification (default: None).
-            Applies to both model and batches.
+            Applies to batches and the model after starting the engine.
         non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
@@ -108,9 +115,6 @@ def create_supervised_evaluator(
     """
     metrics = metrics or {}
 
-    if device:
-        model.to(device)
-
     def _inference(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
         model.eval()
         with torch.no_grad():
@@ -118,9 +122,15 @@ def _inference(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tupl
             y_pred = model(x)
             return output_transform(x, y, y_pred)
 
-    engine = Engine(_inference)
+    evaluator = Engine(_inference)
+
+    if device is not None:
+
+        @evaluator.on(Events.STARTED)
+        def move_model(engine):
+            model.to(device)
 
     for name, metric in metrics.items():
-        metric.attach(engine, name)
+        metric.attach(evaluator, name)
 
-    return engine
+    return evaluator


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -7,6 +7,7 @@
 from torch.optim import SGD
 
 from ignite.engine import create_supervised_trainer, create_supervised_evaluator
+from ignite.engine.engine import Events
 from ignite.metrics import MeanSquaredError
 
 
@@ -100,7 +101,24 @@ def test_create_supervised_trainer_on_cuda():
     assert model.bias.item() == approx(0.8)
 
 
-def test_create_supervised():
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_create_supervised_trainer_on_cuda_with_model_on_cpu_after_init():
+    model = Linear(1, 1)
+    model.weight.data.zero_()
+    model.bias.data.zero_()
+    optimizer = SGD(model.parameters(), 0.1)
+
+    trainer = create_supervised_trainer(model, optimizer, mse_loss, device="cuda")
+
+    x = torch.tensor([[1.0], [2.0]])
+    y = torch.tensor([[3.0], [5.0]])
+    data = [(x, y)]
+
+    model.to("cpu")
+    trainer.run(data)
+
+
+def test_create_supervised_evaluator():
     model = Linear(1, 1)
     model.weight.data.zero_()
     model.bias.data.zero_()
@@ -123,7 +141,7 @@ def test_create_supervised():
     assert model.bias.item() == approx(0.0)
 
 
-def test_create_supervised_on_cpu():
+def test_create_supervised_evaluator_on_cpu():
     model = Linear(1, 1)
     model.weight.data.zero_()
     model.bias.data.zero_()
@@ -173,7 +191,7 @@ def test_create_supervised_evaluator_traced_on_cpu():
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
-def test_create_supervised_on_cuda():
+def test_create_supervised_evaluator_on_cuda():
     model = Linear(1, 1)
     model.weight.data.zero_()
     model.bias.data.zero_()
@@ -196,7 +214,23 @@ def test_create_supervised_on_cuda():
     assert model.bias.item() == approx(0.0)
 
 
-def test_create_supervised_with_metrics():
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_create_supervised_evaluator_on_cuda_with_model_on_cpu_after_init():
+    model = Linear(1, 1)
+    model.weight.data.zero_()
+    model.bias.data.zero_()
+
+    evaluator = create_supervised_evaluator(model, device="cuda")
+
+    x = torch.tensor([[1.0], [2.0]])
+    y = torch.tensor([[3.0], [5.0]])
+    data = [(x, y)]
+
+    model.to("cpu")
+    evaluator.run(data)
+
+
+def test_create_supervised_evaluator_with_metrics():
     model = Linear(1, 1)
     model.weight.data.zero_()
     model.bias.data.zero_()