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