| 1268 | return handle |
| 1269 | |
| 1270 | def _apply(self, fn): |
| 1271 | if not hasattr(self, "cpg"): |
| 1272 | self.cpg = None |
| 1273 | if self.cpg is not None: |
| 1274 | self.cpg = None |
| 1275 | warnings.warn( |
| 1276 | "deleted ContiguousParamsGroup since creating it before " |
| 1277 | "apply operations like to(), to_global() will cause error." |
| 1278 | ) |
| 1279 | |
| 1280 | # A dict to store tensors that has already been applied. |
| 1281 | # There is no need to apply multiple times on a same tensor. |
| 1282 | if self._oneflow_internal_module_tensor_applied_dict__ is None: |
| 1283 | self._oneflow_internal_module_tensor_applied_dict__ = dict() |
| 1284 | |
| 1285 | for module in self.children(): |
| 1286 | module._oneflow_internal_module_tensor_applied_dict__ = ( |
| 1287 | self._oneflow_internal_module_tensor_applied_dict__ |
| 1288 | ) |
| 1289 | module._apply(fn) |
| 1290 | module._oneflow_internal_module_tensor_applied_dict__ = None |
| 1291 | |
| 1292 | def can_use_assign_copy(tensor, tensor_applied): |
| 1293 | return tensor.is_local == tensor_applied.is_local |
| 1294 | |
| 1295 | for (key, param) in self._parameters.items(): |
| 1296 | if param is None: |
| 1297 | continue |
| 1298 | |
| 1299 | need_apply = False |
| 1300 | if param not in self._oneflow_internal_module_tensor_applied_dict__: |
| 1301 | need_apply = True |
| 1302 | assert isinstance(param, Parameter) |
| 1303 | assert param.is_leaf |
| 1304 | with flow.no_grad(): |
| 1305 | param_applied = fn(param) |
| 1306 | param_applied.requires_grad = param.requires_grad |
| 1307 | |
| 1308 | if param.grad is not None: |
| 1309 | assert param.grad.is_leaf |
| 1310 | with flow.no_grad(): |
| 1311 | grad_applied = fn(param.grad) |
| 1312 | grad_applied.requires_grad = param.grad.requires_grad |
| 1313 | param_applied.grad = grad_applied |
| 1314 | else: |
| 1315 | param_applied = self._oneflow_internal_module_tensor_applied_dict__[ |
| 1316 | param |
| 1317 | ] |
| 1318 | |
| 1319 | if can_use_assign_copy(param_applied, param): |
| 1320 | if need_apply: |
| 1321 | self._parameters[key].data = param_applied |
| 1322 | self._oneflow_internal_module_tensor_applied_dict__[ |
| 1323 | param |
| 1324 | ] = param_applied |
| 1325 | else: |
| 1326 | # The parameter's data has already been set when it can use assign copy. |
| 1327 | pass |