已合并
refactor: remove LSTM.forward monkey-patch #43893
ymdsall创建于 8月6日
refactor: remove LSTM.forward monkey-patch #43893
已合并
共 1 个文件变更+0-139
| @@ -170,144 +170,6 @@ def cast_weight(self, device): | |||
| 170 | sub_module.cast_weight(device) | 170 | sub_module.cast_weight(device) |
| 171 | 171 | ||
| 172 | 172 | ||
| 173 | -def _lstm_forward(self, input1, hx=None): | ||
| 174 | - self._update_flat_weights() | ||
| 175 | - | ||
| 176 | - orig_input = input1 | ||
| 177 | - batch_sizes = None | ||
| 178 | - num_directions = 2 if self.bidirectional else 1 | ||
| 179 | - real_hidden_size = self.proj_size if self.proj_size > 0 else self.hidden_size | ||
| 180 | - if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence): | ||
| 181 | - input1, batch_sizes, sorted_indices, unsorted_indices = input1 | ||
| 182 | - max_batch_size = batch_sizes[0] | ||
| 183 | - max_batch_size = int(max_batch_size) | ||
| 184 | - if hx is None: | ||
| 185 | - h_zeros = torch.zeros( | ||
| 186 | - self.num_layers * num_directions, | ||
| 187 | - max_batch_size, | ||
| 188 | - real_hidden_size, | ||
| 189 | - dtype=input1.dtype, | ||
| 190 | - device=input1.device, | ||
| 191 | - ) | ||
| 192 | - c_zeros = torch.zeros( | ||
| 193 | - self.num_layers * num_directions, | ||
| 194 | - max_batch_size, | ||
| 195 | - self.hidden_size, | ||
| 196 | - dtype=input1.dtype, | ||
| 197 | - device=input1.device, | ||
| 198 | - ) | ||
| 199 | - hx = (h_zeros, c_zeros) | ||
| 200 | - else: | ||
| 201 | - hx = self.permute_hidden(hx, sorted_indices) | ||
| 202 | - else: | ||
| 203 | - if input1.dim() not in (2, 3): | ||
| 204 | - raise ValueError( | ||
| 205 | - f"LSTM: Expected input to be 2D or 3D, got {input1.dim()}D instead" | ||
| 206 | - ) | ||
| 207 | - is_batched = input1.dim() == 3 | ||
| 208 | - batch_dim = 0 if self.batch_first else 1 | ||
| 209 | - if not is_batched: | ||
| 210 | - input1 = input1.unsqueeze(batch_dim) | ||
| 211 | - max_batch_size = input1.size(0) if self.batch_first else input1.size(1) | ||
| 212 | - sorted_indices = None | ||
| 213 | - unsorted_indices = None | ||
| 214 | - if hx is None: | ||
| 215 | - h_zeros = torch.zeros( | ||
| 216 | - self.num_layers * num_directions, | ||
| 217 | - max_batch_size, | ||
| 218 | - real_hidden_size, | ||
| 219 | - dtype=input1.dtype, | ||
| 220 | - device=input1.device, | ||
| 221 | - ) | ||
| 222 | - c_zeros = torch.zeros( | ||
| 223 | - self.num_layers * num_directions, | ||
| 224 | - max_batch_size, | ||
| 225 | - self.hidden_size, | ||
| 226 | - dtype=input1.dtype, | ||
| 227 | - device=input1.device, | ||
| 228 | - ) | ||
| 229 | - hx = (h_zeros, c_zeros) | ||
| 230 | - self.check_forward_args(input1, hx, batch_sizes) | ||
| 231 | - else: | ||
| 232 | - if is_batched: | ||
| 233 | - if hx[0].dim() != 3 or hx[1].dim() != 3: | ||
| 234 | - msg = ( | ||
| 235 | - "For batched 3-D input, hx and cx should " | ||
| 236 | - f"also be 3-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors" | ||
| 237 | - ) | ||
| 238 | - raise RuntimeError(msg) | ||
| 239 | - else: | ||
| 240 | - if hx[0].dim() != 2 or hx[1].dim() != 2: | ||
| 241 | - msg = ( | ||
| 242 | - "For unbatched 2-D input, hx and cx should " | ||
| 243 | - f"also be 2-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors" | ||
| 244 | - ) | ||
| 245 | - raise RuntimeError(msg) | ||
| 246 | - hx = (hx[0].unsqueeze(1), hx[1].unsqueeze(1)) | ||
| 247 | - self.check_forward_args(input1, hx, batch_sizes) | ||
| 248 | - hx = self.permute_hidden(hx, sorted_indices) | ||
| 249 | - | ||
| 250 | - if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence): | ||
| 251 | - self.check_forward_args(input1, hx, batch_sizes) | ||
| 252 | - if batch_sizes is None: | ||
| 253 | - result = torch._VF.lstm( | ||
| 254 | - input1, | ||
| 255 | - hx, | ||
| 256 | - self._flat_weights, | ||
| 257 | - self.bias, | ||
| 258 | - self.num_layers, | ||
| 259 | - self.dropout, | ||
| 260 | - self.training, | ||
| 261 | - self.bidirectional, | ||
| 262 | - self.batch_first, | ||
| 263 | - ) | ||
| 264 | - else: | ||
| 265 | - if batch_sizes.device != input1.device: | ||
| 266 | - batch_sizes_npu = batch_sizes.to(input1.device) | ||
| 267 | - result_tmp = torch._VF.lstm( | ||
| 268 | - input1, | ||
| 269 | - batch_sizes_npu, | ||
| 270 | - hx, | ||
| 271 | - self._flat_weights, | ||
| 272 | - self.bias, | ||
| 273 | - self.num_layers, | ||
| 274 | - self.dropout, | ||
| 275 | - self.training, | ||
| 276 | - self.bidirectional, | ||
| 277 | - ) | ||
| 278 | - # pack-lstm-pad时，保持有效T0时序内pad进行lstm定长计算，输出为pack且shape转换[T0*B, *] | ||
| 279 | - if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence): | ||
| 280 | - shape = [result_tmp[0].shape[0] * result_tmp[0].shape[1]] | ||
| 281 | - if result_tmp[0].dim() > 2: | ||
| 282 | - shape = shape + list(result_tmp[0].shape[2:]) | ||
| 283 | - result = (result_tmp[0].reshape(shape),) + result_tmp[1:] | ||
| 284 | - else: | ||
| 285 | - result = torch._VF.lstm( | ||
| 286 | - input1, | ||
| 287 | - batch_sizes, | ||
| 288 | - hx, | ||
| 289 | - self._flat_weights, | ||
| 290 | - self.bias, | ||
| 291 | - self.num_layers, | ||
| 292 | - self.dropout, | ||
| 293 | - self.training, | ||
| 294 | - self.bidirectional, | ||
| 295 | - ) | ||
| 296 | - output = result[0] | ||
| 297 | - hidden = result[1:] | ||
| 298 | - | ||
| 299 | - if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence): | ||
| 300 | - output_packed = torch.nn.utils.rnn.PackedSequence( | ||
| 301 | - output, batch_sizes, sorted_indices, unsorted_indices | ||
| 302 | - ) | ||
| 303 | - return output_packed, self.permute_hidden(hidden, unsorted_indices) | ||
| 304 | - else: | ||
| 305 | - if not is_batched: | ||
| 306 | - output = output.squeeze(batch_dim) | ||
| 307 | - hidden = (hidden[0].squeeze(1), hidden[1].squeeze(1)) | ||
| 308 | - return output, self.permute_hidden(hidden, unsorted_indices) | ||
| 309 | - | ||
| 310 | - | ||
| 311 | def _ddp_init_helper( | 173 | def _ddp_init_helper( |
| 312 | self, | 174 | self, |
| 313 | parameters, | 175 | parameters, |
| @@ -422,7 +284,6 @@ def _apply_module_patch(): | |||
| 422 | torch.nn.Module.npu = npu | 284 | torch.nn.Module.npu = npu |
| 423 | torch.nn.Module.to = to | 285 | torch.nn.Module.to = to |
| 424 | torch.nn.Module.cast_weight = cast_weight | 286 | torch.nn.Module.cast_weight = cast_weight |
| 425 | - torch.nn.modules.rnn.LSTM.forward = _lstm_forward | ||
| 426 | torch.utils.data.dataloader._MultiProcessingDataLoaderIter.__init__ = ( | 287 | torch.utils.data.dataloader._MultiProcessingDataLoaderIter.__init__ = ( |
| 427 | _mpdl_iter_init | 288 | _mpdl_iter_init |
| 428 | ) | 289 | ) |