[
https://issues.apache.org/jira/browse/BEAM-14368?page=com.atlassian.jira.plugin.system.issuetabpanels:all-tabpanel
]
Anand Inguva updated BEAM-14368:
--------------------------------
Description:
Loading pytorch model as whole has some issues with pickling. Investigate it
with running some experiments. If the model size is too large, the current
implementation of the RunInference for PyTorch would fail because of memory
limits.
1. We can pass the model class to the `load_model` of PyTorchModelLoader and
load the model there. This wouldn't pickle the model object but would pickle
the class and the model would be instantiated on the workers.
was:Loading pytorch model as whole has some issues with pickling. Investigate
it with running some experiments. If the model size is too large, the current
implementation of the RunInference for PyTorch would fail because of memory
limits.
> Investigate load state_dict vs loading whole model
> --------------------------------------------------
>
> Key: BEAM-14368
> URL: https://issues.apache.org/jira/browse/BEAM-14368
> Project: Beam
> Issue Type: Sub-task
> Components: sdk-py-core
> Reporter: Anand Inguva
> Assignee: Anand Inguva
> Priority: P2
> Labels: run-inference
>
> Loading pytorch model as whole has some issues with pickling. Investigate it
> with running some experiments. If the model size is too large, the current
> implementation of the RunInference for PyTorch would fail because of memory
> limits.
>
> 1. We can pass the model class to the `load_model` of PyTorchModelLoader and
> load the model there. This wouldn't pickle the model object but would pickle
> the class and the model would be instantiated on the workers.
--
This message was sent by Atlassian Jira
(v8.20.7#820007)