(self,
init_optimizer,
static_loss_scale=1.0,
dynamic_loss_scale=False,
verbose=False)
| 132 | """ |
| 133 | |
| 134 | def __init__(self, |
| 135 | init_optimizer, |
| 136 | static_loss_scale=1.0, |
| 137 | dynamic_loss_scale=False, |
| 138 | verbose=False): |
| 139 | if not torch.cuda.is_available: |
| 140 | raise SystemError("Cannot use fp16 without CUDA.") |
| 141 | |
| 142 | self.verbose = verbose |
| 143 | |
| 144 | self.optimizer = init_optimizer |
| 145 | # init_state_dict sets up an alternative way to cast per-param state |
| 146 | # tensors. |
| 147 | # Stashing here in case https://github.com/pytorch/pytorch/issues/7733 |
| 148 | # makes it necessary. |
| 149 | # init_state_dict = init_optimizer.state_dict() |
| 150 | |
| 151 | self.fp16_groups = [] |
| 152 | self.fp32_from_fp16_groups = [] |
| 153 | self.fp32_from_fp32_groups = [] |
| 154 | for i, param_group in enumerate(self.optimizer.param_groups): |
| 155 | self.maybe_print("FP16_Optimizer processing param group {}:".format(i)) |
| 156 | fp16_params_this_group = [] |
| 157 | fp32_params_this_group = [] |
| 158 | fp32_from_fp16_params_this_group = [] |
| 159 | for i, param in enumerate(param_group['params']): |
| 160 | if param.requires_grad: |
| 161 | if param.type() == 'torch.cuda.HalfTensor': |
| 162 | self.maybe_print( |
| 163 | "FP16_Optimizer received torch.cuda.HalfTensor with" |
| 164 | " {}".format(param.size())) |
| 165 | fp16_params_this_group.append(param) |
| 166 | master_param = param.detach().clone().float() |
| 167 | master_param.requires_grad = True |
| 168 | param_group['params'][i] = master_param |
| 169 | fp32_from_fp16_params_this_group.append(master_param) |
| 170 | # Reset existing state dict key to the new master param. |
| 171 | # We still need to recast per-param state tensors, |
| 172 | # if any, to FP32. |
| 173 | if param in self.optimizer.state: |
| 174 | self.optimizer.state[master_param] = self.optimizer.state.pop(param) |
| 175 | elif param.type() == 'torch.cuda.FloatTensor': |
| 176 | self.maybe_print("FP16_Optimizer received " |
| 177 | "torch.cuda.FloatTensor with {}".format(param.size())) |
| 178 | fp32_params_this_group.append(param) |
| 179 | param_group['params'][i] = param |
| 180 | else: |
| 181 | raise TypeError("Wrapped parameters must be either " |
| 182 | "torch.cuda.FloatTensor or torch.cuda.HalfTensor. " |
| 183 | "Received {}".format(param.type())) |
| 184 | |
| 185 | self.fp16_groups.append(fp16_params_this_group) |
| 186 | self.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group) |
| 187 | self.fp32_from_fp32_groups.append(fp32_params_this_group) |
| 188 | |
| 189 | # Leverage state_dict() and load_state_dict() to recast preexisting |
| 190 | # per-param state tensors |
| 191 | self.optimizer.load_state_dict(self.optimizer.state_dict()) |
nothing calls this directly
no test coverage detected