(
self,
train_net,
train_init_net,
grad_map,
blob_to_device=None,
)
| 703 | modify_output_record=modify_output_record) |
| 704 | |
| 705 | def apply_optimizers( |
| 706 | self, |
| 707 | train_net, |
| 708 | train_init_net, |
| 709 | grad_map, |
| 710 | blob_to_device=None, |
| 711 | ): |
| 712 | CPU = muji.OnCPU() |
| 713 | # if given, blob_to_device is a map from blob to device_option |
| 714 | blob_to_device = blob_to_device or {} |
| 715 | for param, optimizer in self.param_to_optim.items(): |
| 716 | assert optimizer is not None, \ |
| 717 | "default optimizer must have been set in add_layer" |
| 718 | # note that not all params has gradient and thus we sent None if |
| 719 | # gradient does not exists |
| 720 | device = get_param_device( |
| 721 | param, |
| 722 | grad_map.get(str(param)), |
| 723 | param_to_device=blob_to_device, |
| 724 | default_device=CPU, |
| 725 | ) |
| 726 | if device is not None: |
| 727 | # extra info is not applicable for optimizers |
| 728 | del device.extra_info[:] |
| 729 | |
| 730 | with core.DeviceScope(device): |
| 731 | optimizer( |
| 732 | train_net, train_init_net, param, grad_map.get(str(param))) |
| 733 | |
| 734 | def _GetOne(self): |
| 735 | return self.global_constants['ONE'] |
no test coverage detected