def test_compute_losses(sketches, node_feature_mapping, edge_feature_mapping):
sequences = list(map(graph_dataset.sketch_to_sequence, sketches))
dataset = graph_dataset.GraphDataset(sequences,
node_feature_mapping,
edge_feature_mapping,
seed=36)
batch = [dataset[i] for i in range(10)]
batch_input = graph_dataset.collate(batch, node_feature_mapping,
edge_feature_mapping)
feature_dimensions = {
**node_feature_mapping.feature_dimensions,
**edge_feature_mapping.feature_dimensions
}
model = graph_model.make_graph_model(32, feature_dimensions)
model_output = model(batch_input)
losses, _, edge_metrics, node_metrics = graph_model.compute_losses(
model_output, batch_input, feature_dimensions)
assert isinstance(losses, dict)
avg_losses = graph_model.compute_average_losses(
losses, batch_input['graph_counts'])
assert isinstance(avg_losses, dict)
for t in edge_metrics:
assert edge_metrics[t][0].shape == edge_metrics[t][1].shape
for t in edge_metrics:
assert node_metrics[t][0].shape == node_metrics[t][1].shape
def compute_likelihood(self, seqs):
for i, seq in enumerate(seqs):
batch = batch_from_example(seq, self.node_feature_mapping,
self.edge_feature_mapping)
batch_device = training.load_cuda_async(batch, device=self.device)
with torch.no_grad():
readout = self.model(batch_device)
losses, *_ = graph_model.compute_losses(
readout, batch_device, self._feature_dimensions)
loss = _total_loss(losses).cpu().item()
yield loss, len(seq)
def single_step(self, batch, global_step):
self.opt.zero_grad()
profile_path = self._get_profile_path(global_step)
with torch.autograd.profiler.profile(enabled=profile_path is not None,
use_cuda=True) as trace:
with torch.autograd.profiler.record_function("forward"):
readout = self.model(batch)
losses, accuracy, edge_metrics, node_metrics = graph_model.compute_losses(
readout, batch, self.feature_dimension)
total_loss = _total_loss(losses)
if self.model.training:
with torch.autograd.profiler.record_function("backward"):
total_loss.backward()
with torch.autograd.profiler.record_function("opt_update"):
self.opt.step()
if profile_path is not None:
with open(profile_path, 'wb') as f:
pickle.dump(trace, f, pickle.HIGHEST_PROTOCOL)
losses = training.map_structure_flat(losses, _detach)
losses = graph_model.compute_average_losses(losses,
batch['graph_counts'])
avg_loss = total_loss.detach() / float(sum(batch['graph_counts']))
losses['average'] = avg_loss
if self.is_leader():
def _record_classification_summaries(metrics, summaries):
for t, (labels, preds) in metrics.items():
for i, cs in enumerate(summaries[t].values()):
cs.record_statistics(labels[:, i], preds[:, i])
_record_classification_summaries(edge_metrics,
self.edge_feature_summaries)
_record_classification_summaries(node_metrics,
self.node_feature_summaries)
return losses, accuracy
def test_compute_model_subnode(sketches):
sequences = list(map(graph_dataset.sketch_to_sequence, sketches))
dataset = graph_dataset.GraphDataset(sequences, seed=36)
batch = [dataset[i] for i in range(10)]
batch_input = graph_dataset.collate(batch)
model = graph_model.make_graph_model(32,
feature_dimensions={},
readout_entity_features=False,
readout_edge_features=False)
result = model(batch_input)
assert 'node_embedding' in result
losses, _, edge_metrics, node_metrics = graph_model.compute_losses(
result, batch_input, {})
assert isinstance(losses, dict)
avg_losses = graph_model.compute_average_losses(
losses, batch_input['graph_counts'])
assert isinstance(avg_losses, dict)