def test_lstm_saturation_embed_runs():
save_path = TEMP_DIRNAME
# Run 2
timeseries_method = 'last_timestep'
model = torch.nn.Sequential().to(device)
lstm = torch.nn.LSTM(10, 88, 2)
model.add_module('lstm', lstm)
writer = CSVandPlottingWriter(save_path, fontsize=16)
saturation = SaturationTracker(save_path, [writer],
model,
stats=['lsat', 'idim', 'embed'],
timeseries_method=timeseries_method,
device=device)
input = torch.randn(5, 3, 10).to(device)
output, (hn, cn) = model(input)
assert saturation.logs['train-covariance-matrix'][
'lstm'].saved_samples.shape == torch.Size([5, 88])
input = torch.randn(8, 3, 10)
output, (hn, cn) = model(input)
assert saturation.logs['train-covariance-matrix'][
'lstm'].saved_samples.shape == torch.Size([8, 88])
saturation.add_saturations()
saturation.close()
return True
def test_dense_saturation_runs_with_many_writers():
save_path = TEMP_DIRNAME
model = torch.nn.Sequential(torch.nn.Linear(10, 88)).to(device)
writer = CSVandPlottingWriter(save_path,
fontsize=16,
primary_metric='test_accuracy')
writer2 = NPYWriter(save_path)
writer3 = PrintWriter()
sat = SaturationTracker(save_path, [writer, writer2, writer3],
model,
stats=['lsat', 'idim'],
device=device)
test_input = torch.randn(5, 10).to(device)
_ = model(test_input)
sat.add_scalar("test_accuracy", 1.0)
sat.add_saturations()
return True
stats=["lsat", "lsat_eval"])
loss_fn = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
steps_iter = trange(2000, desc='steps', leave=True, position=0)
steps_iter.write("{:^80}".format(
"Regression - TwoLayerNet - Hidden layer size {}".format(h)))
for step in steps_iter:
# training step
model.train()
y_pred = model(x)
loss = loss_fn(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# test step
model.eval()
y_pred = model(x_test)
loss_test = loss_fn(y_pred, y_test)
# update statistics
steps_iter.set_description('loss=%g' % loss.item())
stats.add_scalar("train-loss", loss.item())
stats.add_scalar("test-loss", loss_test.item())
stats.add_saturations()
steps_iter.write('\n')
stats.close()
steps_iter.close()
_, predicted = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total = 0
test_loss = 0
correct = 0
model.eval()
for (images, labels) in tqdm(test_loader):
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += torch.sum((predicted == labels)).item()
test_loss += loss.item()
# add some additional metrics we want to keep track of
tracker.add_scalar("accuracy", correct / total)
tracker.add_scalar("loss", test_loss / total)
# add saturation to the mix
tracker.add_saturations()
tracker.save()
# close the tracker to finish training
tracker.close()