Model Ouput Transform

[david-waterworth]

🚀 Feature

The *_training_step implementations used by create_supervised_trainer expect only a single output from the model that's passed directly to the loss function, i.e.

    x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
    y_pred = model(x)
    loss = loss_fn(y_pred, y)

But I find that commonly models return more than predictions, i.e. they may return logits and embeddings, or logits and loss or perhaps a dictionary

It would be nice if you could pass a transform that would apply to the output from the model in order to extract the predictions in the form required by the loss function, i.e.

    x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
    outputs = model(x)
    y_pred = model_transform(outputs)
    loss = loss_fn(y_pred, y)

I have a current work-around - it required a lambda to handle the loss, and output_transform for each metric, i.e.

loss = functools.partial(F.cross_entropy, label_smoothing=params['train']['label_smoothing'])
trainer = create_supervised_trainer(model, 
                                    optimizer, 
                                    # model output is tuple, need to extract first element
                                    loss_fn=lambda output, target : loss(output[0], target), 
                                    prepare_batch=prepare_batch, 
                                    device=params['train']['device'])

evaluator = create_supervised_evaluator(model,
                                        metrics={
                                            # model output is tuple, need to extract first element
                                            'accuracy': Accuracy(output_transform=lambda output : (output[0][0], output[1])),
                                            'loss': Loss(loss, output_transform=lambda output: (output[0][0], output[1]))},
                                        prepare_batch=prepare_batch,
                                        device=params['train']['device'])

It works fine but it was a bit fiddly to make sure I'd done it correctly, it would be a lot easier if I could just pass a single function to create_supervised_trainer i.e. lambda output: output[0]

[vfdev-5]

@david-waterworth yes, your suggestion makes sense, thanks ! If you would like to send PR you are more than welcome :)

By the way, create_supervised_* are just helper functions showing how to create an engine. Basically, we suggest to create custom train_step and pass it directly to Engine in order to keep flexibility:

device = "cuda"
non_blocking = True

def train_step(engine, batch):
   x, y = batch[0].to(device=device, non_blocking=non_blocking), batch[1].to(device=device, non_blocking=non_blocking)
   outputs = model(x)  # E.g. model outputs a dictionary with logits and loss

   # ... backward pass ...   

   return {
       "loss": outputs["loss"],
       "logits": outputs["logits"]
   }

trainer = Engine(train_step)

[guptaaryan16]

Hey @vfdev-5 and @david-waterworth, can I take up this issue

[david-waterworth]

That's fine with me!

[vfdev-5]

@guptaaryan16 yes, go ahead

[guptaaryan16]

Hey @vfdev-5 in the API of supervised trainer should I replace the output_transform with loss_function_transform and make changes accordingly as it seems to make more sense by the naming convention. Also output_transform can be changed to satisfy the problem in the issue

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[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
    deterministic: bool = False,
    amp_mode: Optional[str] = None,
    scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
    gradient_accumulation_steps: int = 1,
)

[vfdev-5]

@guptaaryan16 I think it is about providing model_transform arg as described in #2837 (comment)

    x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
    outputs = model(x)
    y_pred = model_transform(outputs)
    loss = loss_fn(y_pred, y)

[david-waterworth]

That's correct, as a test, assume model has the following signature

def model(x):
    ...
   return logits, embeddings

in this case, model_transform would unpack the tuple, i.e.

def model_transform(outputs):
    return outputs[0]

[guptaaryan16]

Hey thanks for the clarification @vfdev-5 @david-waterworth