Changes
1 changed files (+77/-0)
-
train.py (new)
-
@@ -0,0 +1,77 @@from argparse import ArgumentParser import torch from torch import nn from dataset import Dataset from model import Model parser = ArgumentParser() parser.add_argument('-d', '--device', help='device to train with', default='cpu') parser.add_argument('-i', '--input', help='training data file', default='data') parser.add_argument('-e', '--epochs', help='number of epochs to train for', default=100) parser.add_argument('-s', '--seq-size', help='sequence size', default=32) parser.add_argument('-b', '--batch-size', help='size of each training batch', default=256) args = parser.parse_args() # Prepare dataloader dataset = Dataset(args.input, args.seq_size) dataloader = DataLoader(dataset, batch_size=args.batch_size) print(len(dataloader)) # Prepare model model = Model(dataset, 512, 512, 3, 0.2).to(device) print(model) loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for t in range(args.epochs): model.train() state_h, state_c = net.zero_state(flags.batch_size) state_h = state_h.to(device) state_c = state_c.to(device) iteration = 0 for batch, (X, y) in enumerate(dataloader): iteration += 1 optimizer.zero_grad() x = torch.tensor(x).to(device) y = torch.tensor(y).to(device) # Compute prediction error logits, (state_h, state_c) = net(x, (state_h, state_c)) loss = loss_fn(logits.transpose(1, 2), y) loss_value = loss.item() # Backpropogation loss.backward() state_h = state_h.detach() state_c = state_c.detach() _ = torch.nn.utils.clip_grad_norm_( model.parameters(), flags.gradients_norm) optimizer.step() if iteration % 10 == 0: print('Epoch: {}/{}'.format(t, args.epochs), 'Iteration: {}'.format(iteration), 'Loss: {}'.format(loss_value)) if iteration % 1000 == 0: predict(device, net, flags.initial_words, n_vocab, vocab_to_int, int_to_vocab, top_k=3) torch.save(net.state_dict(), 'checkpoint/model-{}.pth'.format(iteration))
-