def test_state_attributes():
dataloader = [1, 2, 3]
evaluator = Evaluator(MagicMock(return_value=1))
state = evaluator.run(dataloader)
assert state.iteration == 3
assert state.output == 1
assert state.batch == 3
assert state.dataloader == dataloader
def test_current_validation_iteration_counter_increases_every_iteration():
validation_batches = [1, 2, 3]
evaluator = Evaluator(MagicMock(return_value=1))
num_runs = 5
class IterationCounter(object):
def __init__(self):
self.current_iteration_count = 1
self.total_count = 0
def __call__(self, evaluator, state):
assert state.iteration == self.current_iteration_count
self.current_iteration_count += 1
self.total_count += 1
def clear(self):
self.current_iteration_count = 1
iteration_counter = IterationCounter()
def clear_counter(evaluator, state, counter):
counter.clear()
evaluator.add_event_handler(Events.STARTED, clear_counter,
iteration_counter)
evaluator.add_event_handler(Events.ITERATION_STARTED, iteration_counter)
for _ in range(num_runs):
evaluator.run(validation_batches)
assert iteration_counter.total_count == num_runs * len(validation_batches)
def test_terminate_stops_evaluator_when_called_during_iteration():
num_iterations = 10
iteration_to_stop = 3 # i.e. part way through the 3rd validation run
evaluator = Evaluator(MagicMock(return_value=1))
def start_of_iteration_handler(evaluator, state):
if state.iteration == iteration_to_stop:
evaluator.terminate()
evaluator.add_event_handler(Events.ITERATION_STARTED,
start_of_iteration_handler)
state = evaluator.run([None] * num_iterations)
# should complete the iteration when terminate called but not increment counter
assert state.iteration == iteration_to_stop
def test_timer():
sleep_t = 0.2
n_iter = 3
def _train_func(batch):
time.sleep(sleep_t)
def _test_func(batch):
time.sleep(sleep_t)
trainer = Trainer(_train_func)
tester = Evaluator(_test_func)
t_total = Timer()
t_batch = Timer(average=True)
t_train = Timer()
t_total.attach(trainer)
t_batch.attach(trainer,
pause=Events.ITERATION_COMPLETED,
resume=Events.ITERATION_STARTED,
step=Events.ITERATION_COMPLETED)
t_train.attach(trainer,
pause=Events.EPOCH_COMPLETED,
resume=Events.EPOCH_STARTED)
@trainer.on(Events.EPOCH_COMPLETED)
def run_validation(trainer, state):
tester.run(range(n_iter))
# Run "training"
trainer.run(range(n_iter))
def _equal(lhs, rhs):
return round(lhs, 1) == round(rhs, 1)
assert _equal(t_total.value(), (2 * n_iter * sleep_t))
assert _equal(t_batch.value(), (sleep_t))
assert _equal(t_train.value(), (n_iter * sleep_t))
t_total.reset()
assert _equal(t_total.value(), 0.0)
def test_evaluation_iteration_events_are_fired():
evaluator = Evaluator(MagicMock(return_value=1))
mock_manager = Mock()
iteration_started = Mock()
evaluator.add_event_handler(Events.ITERATION_STARTED, iteration_started)
iteration_complete = Mock()
evaluator.add_event_handler(Events.ITERATION_COMPLETED, iteration_complete)
mock_manager.attach_mock(iteration_started, 'iteration_started')
mock_manager.attach_mock(iteration_complete, 'iteration_complete')
batches = [(1, 2), (3, 4), (5, 6)]
evaluator.run(batches)
assert iteration_started.call_count == len(batches)
assert iteration_complete.call_count == len(batches)
expected_calls = []
for i in range(len(batches)):
expected_calls.append(call.iteration_started(evaluator))
expected_calls.append(call.iteration_complete(evaluator))
assert mock_manager.mock_calls == expected_calls
def main(args, logger):
if args.reverse:
preprocessing = reverse_sentence
else:
preprocessing = None
SRC = data.Field(init_token=BOS,
eos_token=EOS,
include_lengths=True,
preprocessing=preprocessing)
TRG = data.Field(init_token=BOS, eos_token=EOS)
if args.dataset == 'enja':
train, val, test = SmallEnJa.splits(exts=('.en', '.ja'),
fields=(SRC, TRG))
elif args.dataset == 'wmt14':
train, val, test = WMT14.splits(
exts=('.en', '.de'),
fields=(SRC, TRG),
filter_pred=lambda ex: len(ex.src) <= 50 and len(ex.trg) <= 50)
SRC.build_vocab(train.src, max_size=args.src_vocab)
TRG.build_vocab(train.trg, max_size=args.trg_vocab)
stoi = TRG.vocab.stoi
train_iter, val_iter = data.BucketIterator.splits(
(train, val),
batch_sizes=(args.batch, args.batch * 2),
repeat=False,
sort_within_batch=True,
sort_key=SmallEnJa.sort_key,
device=args.gpu[0] if len(args.gpu) == 1 else -1)
test_iter = data.Iterator(test,
batch_size=1,
repeat=False,
sort=False,
train=False,
device=args.gpu[0] if args.gpu else -1)
model = Seq2Seq(len(SRC.vocab), args.embed, args.encoder_hidden,
len(TRG.vocab), args.embed, args.decoder_hidden,
args.encoder_layers, not args.encoder_unidirectional,
args.decoder_layers, SRC.vocab.stoi[SRC.pad_token],
stoi[TRG.pad_token], stoi[TRG.init_token],
stoi[TRG.eos_token], args.dropout_ratio,
args.attention_type)
model.prepare_translation(TRG.vocab.itos, args.max_length)
translate = model.translate
if len(args.gpu) >= 2:
model = DataParallel(model, device_ids=args.gpu, dim=1).cuda()
elif len(args.gpu) == 1:
model.cuda(args.gpu[0])
if args.optim == 'adam':
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
scheduler = None
elif args.optim == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=args.learning_rate)
scheduler = MultiStepLR(optimizer,
milestones=list(range(8, 12)),
gamma=0.5)
trainer = Trainer(
TeacherForceUpdater(model, optimizer, model, args.gradient_clipping))
evaluator = Evaluator(TeacherForceInference(model, model))
trainer.add_event_handler(Events.EPOCH_COMPLETED,
log_training_average_nll,
logger=logger)
trainer.add_event_handler(Events.EPOCH_COMPLETED,
Evaluate(evaluator, val_iter, epoch_interval=1))
evaluator.add_event_handler(Events.COMPLETED,
get_log_validation_ppl(val.trg),
logger=logger)
if args.best_file is not None:
evaluator.add_event_handler(Events.COMPLETED,
BestModelSnapshot(model, 'ppl', 1e10, le),
args.best_file, logger)
trainer.add_event_handler(Events.COMPLETED,
ComputeBleu(model, test.trg, translate),
test_iter, args.best_file, logger)
if scheduler is not None:
trainer.add_event_handler(Events.EPOCH_STARTED,
lambda trainer: scheduler.step())
trainer.run(train_iter, max_epochs=args.epoch)
def run():
train_dataloader = data.DataLoader(CMAPSSData(train=True),
shuffle=True,
batch_size=300,
pin_memory=USE_CUDA,
collate_fn=collate_fn)
validation_dataloader = data.DataLoader(CMAPSSData(train=False),
shuffle=True,
batch_size=1000,
pin_memory=USE_CUDA,
collate_fn=collate_fn)
model = TTE(24, 128, n_layers=1, dropout=0.2)
print(model)
if USE_CUDA:
model.cuda()
optimizer = optim.RMSprop(model.parameters(), lr=0.001)
loss_fn = weibull_loglikelihood
def training_update(batch):
model.train()
optimizer.zero_grad()
inputs, lengths, targets = batch
if USE_CUDA:
inputs = inputs.cuda()
targets = targets.cuda()
inputs = Variable(inputs)
targets = Variable(targets)
outputs = model(inputs, lengths)
loss = loss_fn(outputs, targets, lengths)
loss.backward()
optimizer.step()
return loss.data[0]
def validation_inference(batch):
model.eval()
inputs, lengths, targets = batch
if USE_CUDA:
inputs = inputs.cuda()
targets = targets.cuda()
inputs = Variable(inputs, volatile=True)
targets = Variable(targets, volatile=True)
outputs = model(inputs, lengths)
loss = loss_fn(outputs, targets, lengths)
return loss.data[0], outputs.data[:, :,
0], outputs.data[:, :,
1], targets.data
trainer = Trainer(training_update)
evaluator = Evaluator(validation_inference)
progress = Progress()
plot_interval = 1
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
Evaluate(evaluator, validation_dataloader, epoch_interval=1))
@trainer.on(Events.EPOCH_STARTED)
def epoch_started(trainer):
print('Epoch {:4}/{}'.format(trainer.current_epoch,
trainer.max_epochs),
end='')
@trainer.on(Events.ITERATION_COMPLETED)
def iteration_completed(trainer):
if trainer.current_iteration % plot_interval == 0:
avg_loss = trainer.history.simple_moving_average(window_size=100)
values = [('iter', trainer.current_iteration), ('loss', avg_loss)]
progress.update(values)
@evaluator.on(Events.COMPLETED)
def epoch_completed(evaluator):
history = evaluator.history[0]
loss = history[0]
alpha = history[1]
beta = history[2]
target = history[3]
mae = torch.mean(torch.abs(alpha - target[:, :, 0]))
alpha = alpha.mean()
beta = beta.mean()
values = [('val_loss', loss), ('mae', mae), ('alpha', alpha),
('beta', beta)]
progress.update(values, end=True)
trainer.run(train_dataloader, max_epochs=600)
return model
def run(batch_size, val_batch_size, epochs, lr, momentum, log_interval,
logger):
vis = visdom.Visdom()
if not vis.check_connection():
raise RuntimeError(
"Visdom server not running. Please run python -m visdom.server")
data_transform = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_loader = DataLoader(MNIST(download=True,
root=".",
transform=data_transform,
train=True),
batch_size=batch_size,
shuffle=True)
val_loader = DataLoader(MNIST(download=False,
root=".",
transform=data_transform,
train=False),
batch_size=val_batch_size,
shuffle=False)
model = Net()
optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
def training_update_function(batch):
model.train()
optimizer.zero_grad()
data, target = Variable(batch[0]), Variable(batch[1])
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
return loss.data[0]
def validation_inference_function(batch):
model.eval()
data, target = Variable(batch[0], volatile=True), Variable(batch[1])
output = model(data)
loss = F.nll_loss(output, target, size_average=False).data[0]
pred = output.data.max(1, keepdim=True)[1]
correct = pred.eq(target.data.view_as(pred)).sum()
return loss, correct
trainer = Trainer(training_update_function)
evaluator = Evaluator(validation_inference_function)
# trainer event handlers
trainer.add_event_handler(Events.ITERATION_COMPLETED,
log_simple_moving_average,
window_size=100,
metric_name="NLL",
should_log=lambda trainer: trainer.
current_iteration % log_interval == 0,
logger=logger)
trainer.add_event_handler(
Events.ITERATION_COMPLETED,
get_plot_training_loss_handler(vis, plot_every=log_interval))
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
Evaluate(evaluator, val_loader, epoch_interval=1))
# evaluator event handlers
evaluator.add_event_handler(
Events.COMPLETED, get_log_validation_loss_and_accuracy_handler(logger))
evaluator.add_event_handler(Events.COMPLETED,
get_plot_validation_accuracy_handler(vis),
trainer)
# kick everything off
trainer.run(train_loader, max_epochs=epochs)
def test_returns_state():
evaluator = Evaluator(MagicMock(return_value=1))
state = evaluator.run([])
assert isinstance(state, State)
def run(batch_size, val_batch_size, epochs, lr, momentum, log_interval, logger,
log_dir):
writer = SummaryWriter(log_dir=log_dir)
data_transform = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))])
train_loader = DataLoader(MNIST(download=True,
root=".",
transform=data_transform,
train=True),
batch_size=batch_size,
shuffle=True)
val_loader = DataLoader(MNIST(download=False,
root=".",
transform=data_transform,
train=False),
batch_size=val_batch_size,
shuffle=False)
model = Net()
optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
try:
dummy_input = Variable(torch.rand(10, 1, 28, 28))
torch.onnx.export(model, dummy_input, "model.proto", verbose=True)
writer.add_graph_onnx("model.proto")
except ImportError:
pass
def training_update_function(batch):
model.train()
optimizer.zero_grad()
data, target = Variable(batch[0]), Variable(batch[1])
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
return loss.data[0]
def validation_inference_function(batch):
model.eval()
data, target = Variable(batch[0]), Variable(batch[1])
output = model(data)
loss = F.nll_loss(output, target, size_average=False).data[0]
pred = output.data.max(1, keepdim=True)[1]
correct = pred.eq(target.data.view_as(pred)).sum()
return loss, correct
trainer = Trainer(training_update_function)
evaluator = Evaluator(validation_inference_function)
# trainer event handlers
trainer.add_event_handler(Events.ITERATION_COMPLETED,
log_simple_moving_average,
window_size=100,
metric_name="NLL",
should_log=lambda trainer: trainer.
current_iteration % log_interval == 0,
logger=logger)
trainer.add_event_handler(
Events.ITERATION_COMPLETED,
get_plot_training_loss_handler(writer, plot_every=log_interval))
trainer.add_event_handler(
Events.EPOCH_COMPLETED,
Evaluate(evaluator, val_loader, epoch_interval=1))
# evaluator event handlers
evaluator.add_event_handler(
Events.COMPLETED, get_log_validation_loss_and_accuracy_handler(logger))
evaluator.add_event_handler(Events.COMPLETED,
get_plot_validation_loss_handler(writer),
trainer)
evaluator.add_event_handler(Events.COMPLETED,
get_plot_validation_accuracy_handler(writer),
trainer)
# kick everything off
trainer.run(train_loader, max_epochs=epochs)
writer.close()