PyTorchを使ったTransformerによる翻訳の際に長文の文章で学習を行いたい

前提

PyTorchを使ってTransformerによる翻訳モデルを作って翻訳を行いたい．こちらのサイト
(https://www.dskomei.com/entry/2021/05/24/165158)
のコード利用させていただき実際に試したところ，データとしてこちら(http://www.manythings.org/anki/)
の英日のデータセットを用いて学習するとうまくいった．
しかし，長文を学習したく500単語ほどの文章を一つのペアとして学習を行ったところ学習データが50を超えるとRuntimeErrorが出てしまった．

実現したいこと

500単語程度の文章を一つのペアとして大量のデータを用いて学習を行いたい．
そもそも500単語程度の文章(長文)を学習するのは現実的ではないのか，モデルの設計の部分で何か間違っているのか教えていただきたい．

発生している問題・エラーメッセージ

RuntimeError: CUDA out of memory. Tried to allocate 2.27 GiB (GPU 0; 39.59 GiB total capacity; 37.52 GiB already allocated; 250.19 MiB free; 37.82 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

該当のソースコード

Python
1#省略
2
3#以下翻訳モデルの設計
4class Seq2SeqTransformer(nn.Module):
5    
6    def __init__(
7        self, num_encoder_layers: int, num_decoder_layers: int,
8        embedding_size: int, vocab_size_src: int, vocab_size_tgt: int,
9        dim_feedforward:int = 512, dropout:float = 0.1, nhead:int = 8
10    ):
11        
12        super(Seq2SeqTransformer, self).__init__()
13
14        self.token_embedding_src = TokenEmbedding(vocab_size_src, embedding_size)
15        self.positional_encoding = PositionalEncoding(embedding_size, dropout=dropout)
16        encoder_layer = TransformerEncoderLayer(
17            d_model=embedding_size, nhead=nhead, dim_feedforward=dim_feedforward
18        )
19        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
20        
21        self.token_embedding_tgt = TokenEmbedding(vocab_size_tgt, embedding_size)
22        decoder_layer = TransformerDecoderLayer(
23            d_model=embedding_size, nhead=nhead, dim_feedforward=dim_feedforward
24        )
25        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
26        
27        self.output = nn.Linear(embedding_size, vocab_size_tgt)
28
29    def forward(
30        self, src: Tensor, tgt: Tensor,
31        mask_src: Tensor, mask_tgt: Tensor,
32        padding_mask_src: Tensor, padding_mask_tgt: Tensor,
33        memory_key_padding_mask: Tensor
34    ):
35        
36        embedding_src = self.positional_encoding(self.token_embedding_src(src))
37        memory = self.transformer_encoder(embedding_src, mask_src, padding_mask_src)
38        embedding_tgt = self.positional_encoding(self.token_embedding_tgt(tgt))
39        outs = self.transformer_decoder(
40            embedding_tgt, memory, mask_tgt, None,
41            padding_mask_tgt, memory_key_padding_mask
42        )
43        return self.output(outs)
44
45    def encode(self, src: Tensor, mask_src: Tensor):
46        return self.transformer_encoder(self.positional_encoding(self.token_embedding_src(src)), mask_src)
47
48    def decode(self, tgt: Tensor, memory: Tensor, mask_tgt: Tensor):
49        return self.transformer_decoder(self.positional_encoding(self.token_embedding_tgt(tgt)), memory, mask_tgt)
50
51class TokenEmbedding(nn.Module):
52    
53    def __init__(self, vocab_size, embedding_size):
54        
55        super(TokenEmbedding, self).__init__()
56        self.embedding = nn.Embedding(vocab_size, embedding_size)
57        self.embedding_size = embedding_size
58        
59    def forward(self, tokens: Tensor):
60        return self.embedding(tokens.long()) * math.sqrt(self.embedding_size)
61
62class PositionalEncoding(nn.Module):
63    
64    def __init__(self, embedding_size: int, dropout: float, maxlen: int = 5000):
65        super(PositionalEncoding, self).__init__()
66        
67        den = torch.exp(-torch.arange(0, embedding_size, 2) * math.log(10000) / embedding_size)
68        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
69        embedding_pos = torch.zeros((maxlen, embedding_size))
70        embedding_pos[:, 0::2] = torch.sin(pos * den)
71        embedding_pos[:, 1::2] = torch.cos(pos * den)
72        embedding_pos = embedding_pos.unsqueeze(-2)
73
74        self.dropout = nn.Dropout(dropout)
75        self.register_buffer('embedding_pos', embedding_pos)
76
77    def forward(self, token_embedding: Tensor):
78        return self.dropout(token_embedding + self.embedding_pos[: token_embedding.size(0), :])
79
80def create_mask(src, tgt, PAD_IDX):
81    
82    seq_len_src = src.shape[0]
83    seq_len_tgt = tgt.shape[0]
84
85    mask_tgt = generate_square_subsequent_mask(seq_len_tgt, PAD_IDX)
86    mask_src = torch.zeros((seq_len_src, seq_len_src), device=device).type(torch.bool)
87
88    padding_mask_src = (src == PAD_IDX).transpose(0, 1)
89    padding_mask_tgt = (tgt == PAD_IDX).transpose(0, 1)
90    
91    return mask_src, mask_tgt, padding_mask_src, padding_mask_tgt
92
93
94def generate_square_subsequent_mask(seq_len, PAD_IDX):
95    mask = (torch.triu(torch.ones((seq_len, seq_len), device=device)) == 1).transpose(0, 1)
96    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == PAD_IDX, float(0.0))
97    return mask
98
99def train(model, data, optimizer, criterion, PAD_IDX):
100    
101    model.train()
102    losses = 0
103    for src, tgt in tqdm(data):
104        
105        src = src.to(device)
106        tgt = tgt.to(device)
107
108        input_tgt = tgt[:-1, :]
109
110        mask_src, mask_tgt, padding_mask_src, padding_mask_tgt = create_mask(src, input_tgt, PAD_IDX)
111
112        logits = model(
113            src=src, tgt=input_tgt,
114            mask_src=mask_src, mask_tgt=mask_tgt,
115            padding_mask_src=padding_mask_src, padding_mask_tgt=padding_mask_tgt,
116            memory_key_padding_mask=padding_mask_src
117        )
118
119        optimizer.zero_grad()
120
121        output_tgt = tgt[1:, :]
122        loss = criterion(logits.reshape(-1, logits.shape[-1]), output_tgt.reshape(-1))
123        loss.backward()
124
125        optimizer.step()
126        losses += loss.item()
127        
128    return losses / len(data)
129
130
131def evaluate(model, data, criterion, PAD_IDX):
132    
133    model.eval()
134    losses = 0
135    for src, tgt in data:
136        
137        src = src.to(device)
138        tgt = tgt.to(device)
139
140        input_tgt = tgt[:-1, :]
141
142        mask_src, mask_tgt, padding_mask_src, padding_mask_tgt = create_mask(src, input_tgt, PAD_IDX)
143
144        logits = model(
145            src=src, tgt=input_tgt,
146            mask_src=mask_src, mask_tgt=mask_tgt,
147            padding_mask_src=padding_mask_src, padding_mask_tgt=padding_mask_tgt,
148            memory_key_padding_mask=padding_mask_src
149        )
150        
151        output_tgt = tgt[1:, :]
152        loss = criterion(logits.reshape(-1, logits.shape[-1]), output_tgt.reshape(-1))
153        losses += loss.item()
154        
155    return losses / len(data)
156
157vocab_size_src = len(vocab_src)
158vocab_size_tgt = len(vocab_tgt)
159embedding_size = 240
160nhead = 8
161dim_feedforward = 100
162num_encoder_layers = 2
163num_decoder_layers = 2
164dropout = 0.1
165
166model = Seq2SeqTransformer(
167    num_encoder_layers=num_encoder_layers,
168    num_decoder_layers=num_decoder_layers,
169    embedding_size=embedding_size,
170    vocab_size_src=vocab_size_src, vocab_size_tgt=vocab_size_tgt,
171    dim_feedforward=dim_feedforward,
172    dropout=dropout, nhead=nhead
173)
174
175for p in model.parameters():
176    if p.dim() > 1:
177        nn.init.xavier_uniform_(p)
178
179model = model.to(device)
180
181criterion = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)
182
183optimizer = torch.optim.Adam(model.parameters())
184
185#以下のモデルの学習のところでエラー
186epoch = 50
187best_loss = float('Inf')
188best_model = None
189patience = 10
190counter = 0
191
192for loop in range(1, epoch + 1):
193    
194    start_time = time.time()
195    
196    loss_train = train(
197        model=model, data=train_iter, optimizer=optimizer,
198        criterion=criterion, PAD_IDX=PAD_IDX
199    )
200    
201    elapsed_time = time.time() - start_time
202    
203    loss_valid = evaluate(
204        model=model, data=valid_iter, criterion=criterion, PAD_IDX=PAD_IDX
205    )
206    
207    print('[{}/{}] train loss: {:.2f}, valid loss: {:.2f}  [{}{:.0f}s] count: {}, {}'.format(
208        loop, epoch,
209        loss_train, loss_valid,
210        str(int(math.floor(elapsed_time / 60))) + 'm' if math.floor(elapsed_time / 60) > 0 else '',
211        elapsed_time % 60,
212        counter,
213        '**' if best_loss > loss_valid else ''
214    ))
215    
216    if best_loss > loss_valid:
217        best_loss = loss_valid
218        best_model = model
219        counter = 0
220        
221    if counter > patience:
222        break
223    
224    counter += 1
225
226#以下省略