(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder, embedding,
encoder_optimizer, decoder_optimizer, batch_size, clip, max_length=MAX_LENGTH)
| 944 | |
| 945 | |
| 946 | def train(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder, embedding, |
| 947 | encoder_optimizer, decoder_optimizer, batch_size, clip, max_length=MAX_LENGTH): |
| 948 | |
| 949 | # Zero gradients |
| 950 | encoder_optimizer.zero_grad() |
| 951 | decoder_optimizer.zero_grad() |
| 952 | |
| 953 | # Set device options |
| 954 | input_variable = input_variable.to(device) |
| 955 | target_variable = target_variable.to(device) |
| 956 | mask = mask.to(device) |
| 957 | # Lengths for RNN packing should always be on the CPU |
| 958 | lengths = lengths.to("cpu") |
| 959 | |
| 960 | # Initialize variables |
| 961 | loss = 0 |
| 962 | print_losses = [] |
| 963 | n_totals = 0 |
| 964 | |
| 965 | # Forward pass through encoder |
| 966 | encoder_outputs, encoder_hidden = encoder(input_variable, lengths) |
| 967 | |
| 968 | # Create initial decoder input (start with SOS tokens for each sentence) |
| 969 | decoder_input = torch.LongTensor([[SOS_token for _ in range(batch_size)]]) |
| 970 | decoder_input = decoder_input.to(device) |
| 971 | |
| 972 | # Set initial decoder hidden state to the encoder's final hidden state |
| 973 | decoder_hidden = encoder_hidden[:decoder.n_layers] |
| 974 | |
| 975 | # Determine if we are using teacher forcing this iteration |
| 976 | use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False |
| 977 | |
| 978 | # Forward batch of sequences through decoder one time step at a time |
| 979 | if use_teacher_forcing: |
| 980 | for t in range(max_target_len): |
| 981 | decoder_output, decoder_hidden = decoder( |
| 982 | decoder_input, decoder_hidden, encoder_outputs |
| 983 | ) |
| 984 | # Teacher forcing: next input is current target |
| 985 | decoder_input = target_variable[t].view(1, -1) |
| 986 | # Calculate and accumulate loss |
| 987 | mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t]) |
| 988 | loss += mask_loss |
| 989 | print_losses.append(mask_loss.item() * nTotal) |
| 990 | n_totals += nTotal |
| 991 | else: |
| 992 | for t in range(max_target_len): |
| 993 | decoder_output, decoder_hidden = decoder( |
| 994 | decoder_input, decoder_hidden, encoder_outputs |
| 995 | ) |
| 996 | # No teacher forcing: next input is decoder's own current output |
| 997 | _, topi = decoder_output.topk(1) |
| 998 | decoder_input = torch.LongTensor([[topi[i][0] for i in range(batch_size)]]) |
| 999 | decoder_input = decoder_input.to(device) |
| 1000 | # Calculate and accumulate loss |
| 1001 | mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t]) |
| 1002 | loss += mask_loss |
| 1003 | print_losses.append(mask_loss.item() * nTotal) |
no test coverage detected