apache / apache/beam

Investigate load state_dict vs loading whole model

Open
#21,450 4 comments 0 reactions 0 assignees View on GitHub
core ml P3 python run-inference sub-task
Dominant language
Java
Stars
8.7k
Forks
4.7k
Avg merge
1d 20h
Merged PRs (30d)
196

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.

Imported from Jira [BEAM-14368](https://issues.apache.org/jira/browse/BEAM-14368). Original Jira may contain additional context.
Reported by: Anand Inguva.
Subtask of issue #21435

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.