| 610 | # |
| 611 | |
| 612 | def train_epoch(dataloader, encoder, decoder, encoder_optimizer, |
| 613 | decoder_optimizer, criterion): |
| 614 | |
| 615 | total_loss = 0 |
| 616 | for data in dataloader: |
| 617 | input_tensor, target_tensor = data |
| 618 | |
| 619 | encoder_optimizer.zero_grad() |
| 620 | decoder_optimizer.zero_grad() |
| 621 | |
| 622 | encoder_outputs, encoder_hidden = encoder(input_tensor) |
| 623 | decoder_outputs, _, _ = decoder(encoder_outputs, encoder_hidden, target_tensor) |
| 624 | |
| 625 | loss = criterion( |
| 626 | decoder_outputs.view(-1, decoder_outputs.size(-1)), |
| 627 | target_tensor.view(-1) |
| 628 | ) |
| 629 | loss.backward() |
| 630 | |
| 631 | encoder_optimizer.step() |
| 632 | decoder_optimizer.step() |
| 633 | |
| 634 | total_loss += loss.item() |
| 635 | |
| 636 | return total_loss / len(dataloader) |
| 637 | |
| 638 | |
| 639 | ###################################################################### |