Place model on certain device for pytorch models before first inference
(self)
| 156 | raise RuntimeError(info_str) |
| 157 | |
| 158 | def prepare_model(self): |
| 159 | """ Place model on certain device for pytorch models before first inference |
| 160 | """ |
| 161 | self._model_prepare_lock.acquire(timeout=600) |
| 162 | |
| 163 | def _prepare_single(model): |
| 164 | if not isinstance(model, torch.nn.Module) and hasattr( |
| 165 | model, 'model'): |
| 166 | model = model.model |
| 167 | if not isinstance(model, torch.nn.Module): |
| 168 | return |
| 169 | model.eval() |
| 170 | from modelscope.utils.torch_utils import is_on_same_device |
| 171 | if is_on_same_device(model): |
| 172 | model.to(self.device) |
| 173 | |
| 174 | if not self._model_prepare: |
| 175 | # prepare model for pytorch |
| 176 | if self.framework == Frameworks.torch: |
| 177 | if self.has_multiple_models: |
| 178 | for m in self.models: |
| 179 | _prepare_single(m) |
| 180 | if self._compile: |
| 181 | self.models = [ |
| 182 | compile_model(m, **self._compile_options) |
| 183 | for m in self.models |
| 184 | ] |
| 185 | else: |
| 186 | _prepare_single(self.model) |
| 187 | if self._compile: |
| 188 | self.model = compile_model(self.model, |
| 189 | **self._compile_options) |
| 190 | self._model_prepare = True |
| 191 | self._model_prepare_lock.release() |
| 192 | |
| 193 | def _get_framework(self) -> str: |
| 194 | frameworks = [] |
no test coverage detected