yeandy commented on issue #22572:
URL: https://github.com/apache/beam/issues/22572#issuecomment-1212215137
> pytorch calls the model like a callable (which then uses the forward
method IIUC?):
Correct.
---
And thanks @agvdndor for the detailed suggestions!
- For `GenerationModelHandler` I agree that it does not scale well.
- A lambda like `model_inference_fn` could work. The change itself shouldn't
be that hard to implement. However, we need to ask ourselves -- at what point
are we doing too much to address these custom use cases? On the one hand, I
recognize that HuggingFace is very popular, if I'd be remiss if to see a bunch
of potential RunInference users turned away because of how difficult it is to
plug in a HuggingFace model into `PytorchModelHandlerTensor`. On the other
hand, if we can capture 80% of use cases without having this custom infer
function, that might be good enough? If users do require a more tailored
solution, then they probably should be writing up their own `DoFn` anyway
(inspired, of course, by our own implementation). @robertwb What are your
thoughts on adding something like a Generation ModleHandler versus a
`model_inference_fn`?
There are some other workarounds that users could do. Would these be
sufficient solutions to this?
1. Create a wrapper class that inherits from `torch.nn.Module`, and then
override its `forward()` method and calls the model's intended inference
function. (Note: this code is just an example and isn't necessarily the best or
correct way to do this.)
```
class Tacotron2Wrapper(torch.nn.Module):
def __init__(self, model=tacotron2):
super().__init__()
self._model = model
def forward(self, inputs, input_lengths):
mel, _, _ = self._model.infer(inputs, input_lengths)
return mel
```
2. Inherit `ModelHandler`, and change the `run_inference` function to call
`model.infer()` instead of `model()`. This might be easier than the first
solution, but does require the user to copy the other logic correctly.
```
def run_inference(
self,
batch: Sequence[torch.Tensor],
model: torch.nn.Module,
inference_args: Optional[Dict[str, Any]] = None
) -> Iterable[PredictionResult]:
inference_args = {} if not inference_args else inference_args
batched_tensors = torch.stack(batch)
batched_tensors = _convert_to_device(batched_tensors, self._device)
predictions = model.infer(batched_tensors, **inference_args)
return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
