Creates a list of FP32 master parameters for a given model, as in `Training Neural Networks with Mixed Precision: Real Examples`_. Args: model (torch.nn.Module): Existing Pytorch model flat_master (bool, optional, default=False): Flatten the master param
(model, flat_master=False)
| 227 | |
| 228 | |
| 229 | def prep_param_lists(model, flat_master=False): |
| 230 | """ |
| 231 | Creates a list of FP32 master parameters for a given model, as in |
| 232 | `Training Neural Networks with Mixed Precision: Real Examples`_. |
| 233 | |
| 234 | Args: |
| 235 | model (torch.nn.Module): Existing Pytorch model |
| 236 | flat_master (bool, optional, default=False): Flatten the master |
| 237 | parameters into a single tensor, as a performance optimization. |
| 238 | Returns: |
| 239 | A tuple (``model_params``, ``master_params``). ``model_params`` is a |
| 240 | list of the model's parameters for later use with |
| 241 | :func:`model_grads_to_master_grads` and |
| 242 | :func:`master_params_to_model_params`. |
| 243 | ``master_params`` is a list of FP32 master gradients. |
| 244 | If ``flat_master=True``, ``master_params`` will be a list with one |
| 245 | element. |
| 246 | |
| 247 | Example:: |
| 248 | |
| 249 | model_params, master_params = prep_param_lists(model) |
| 250 | |
| 251 | .. warning:: |
| 252 | Currently, if ``flat_master=True``, all the model's parameters must be |
| 253 | the same type. If the model has parameters of different types, use |
| 254 | ``flat_master=False``, or use :class:`FP16_Optimizer`. |
| 255 | |
| 256 | .. _`Training Neural Networks with Mixed Precision: Real Examples`: |
| 257 | http://on-demand.gputechconf.com/gtc/2018/video/S81012/ |
| 258 | """ |
| 259 | model_params = [param for param in model.parameters() if param.requires_grad] |
| 260 | |
| 261 | if flat_master: |
| 262 | # Give the user some more useful error messages |
| 263 | try: |
| 264 | # flatten_dense_tensors returns a contiguous flat array. |
| 265 | # http://pytorch.org/docs/master/_modules/torch/_utils.html |
| 266 | master_params = _flatten_dense_tensors([param.data for param in |
| 267 | model_params]).float() |
| 268 | except: |
| 269 | print("Error in prep_param_lists: model may contain a mixture of parameters " |
| 270 | "of different types. Use flat_master=False, or use F16_Optimizer.") |
| 271 | raise |
| 272 | master_params = torch.nn.Parameter(master_params) |
| 273 | master_params.requires_grad = True |
| 274 | # master_params.register_hook(backwards_debug_hook) |
| 275 | if master_params.grad is None: |
| 276 | master_params.grad = master_params.new(*master_params.size()) |
| 277 | return model_params, [master_params] |
| 278 | else: |
| 279 | master_params = [param.clone().float().detach() for param in model_params] |
| 280 | for param in master_params: |
| 281 | param.requires_grad = True |
| 282 | return model_params, master_params |
| 283 | |
| 284 | |
| 285 | def model_grads_to_master_grads(model_params, master_params, flat_master=False): |