mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
161e931156
Use pyupgrade(https://github.com/asottile/pyupgrade) and flynt to modernize python syntax ```sh pyupgrade --py36-plus --keep-runtime-typing torch/onnx/**/*.py pyupgrade --py36-plus --keep-runtime-typing test/onnx/**/*.py flynt torch/onnx/ --line-length 120 ``` - Use f-strings for string formatting - Use the new `super()` syntax for class initialization - Use dictionary / set comprehension Pull Request resolved: https://github.com/pytorch/pytorch/pull/77935 Approved by: https://github.com/BowenBao
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from torch import nn
|
|
from torch.nn.utils.rnn import PackedSequence
|
|
|
|
|
|
class LstmFlatteningResult(nn.LSTM):
|
|
def forward(self, input, *fargs, **fkwargs):
|
|
output, (hidden, cell) = nn.LSTM.forward(self, input, *fargs, **fkwargs)
|
|
return output, hidden, cell
|
|
|
|
|
|
class LstmFlatteningResultWithSeqLength(nn.Module):
|
|
def __init__(self, input_size, hidden_size, layers, bidirect, dropout, batch_first):
|
|
super().__init__()
|
|
|
|
self.batch_first = batch_first
|
|
self.inner_model = nn.LSTM(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=layers,
|
|
bidirectional=bidirect,
|
|
dropout=dropout,
|
|
batch_first=batch_first,
|
|
)
|
|
|
|
def forward(self, input: PackedSequence, hx=None):
|
|
output, (hidden, cell) = self.inner_model.forward(input, hx)
|
|
return output, hidden, cell
|
|
|
|
|
|
class LstmFlatteningResultWithoutSeqLength(nn.Module):
|
|
def __init__(self, input_size, hidden_size, layers, bidirect, dropout, batch_first):
|
|
super().__init__()
|
|
|
|
self.batch_first = batch_first
|
|
self.inner_model = nn.LSTM(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=layers,
|
|
bidirectional=bidirect,
|
|
dropout=dropout,
|
|
batch_first=batch_first,
|
|
)
|
|
|
|
def forward(self, input, hx=None):
|
|
output, (hidden, cell) = self.inner_model.forward(input, hx)
|
|
return output, hidden, cell
|