Skip to content

rust: support graph sub-plugins #4569

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wchargin merged 7 commits into from
Jan 19, 2021
Merged

rust: support graph sub-plugins #4569

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b53e1a1
rust: support images and generic blob sequences
wchargin Jan 16, 2021
d2ceae4
rust: pass summary metadata to `into_blob_sequence`
wchargin Jan 16, 2021
a1c4362
rust: support audio
wchargin Jan 16, 2021
61a5e01
rust: support tagged run metadata graphs
wchargin Jan 16, 2021
ede7691
rust: support graph sub-plugins
wchargin Jan 16, 2021
22dd050
[rust-graph-subplugins: update diffbase]
wchargin Jan 19, 2021
a4fe0c2
[rust-graph-subplugins: update patch]
wchargin Jan 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 77 additions & 10 deletions tensorboard/data/server/data_compat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ pub(crate) const IMAGES_PLUGIN_NAME: &str = "images";
pub(crate) const AUDIO_PLUGIN_NAME: &str = "audio";
pub(crate) const GRAPHS_PLUGIN_NAME: &str = "graphs";
pub(crate) const GRAPH_TAGGED_RUN_METADATA_PLUGIN_NAME: &str = "graph_tagged_run_metadata";
pub(crate) const GRAPH_RUN_METADATA_PLUGIN_NAME: &str = "graph_run_metadata";
pub(crate) const GRAPH_RUN_METADATA_WITH_GRAPH_PLUGIN_NAME: &str = "graph_run_metadata_graph";
pub(crate) const GRAPH_KERAS_MODEL_PLUGIN_NAME: &str = "graph_keras_model";

/// The inner contents of a single value from an event.
///
Expand Down Expand Up @@ -76,11 +79,17 @@ impl EventValue {

/// Consumes this event value and enriches it into a blob sequence.
///
/// For now, this supports `GraphDef`s, tagged run metadata protos, summaries with `image` or
/// `audio`, or summaries with `tensor` set to a rank-1 tensor of type `DT_STRING`. If the
/// summary metadata indicates that this is audio data, `tensor` may also be a string tensor of
/// shape `[k, 2]`, in which case the second axis is assumed to represent string labels and is
/// dropped entirely.
/// This supports:
///
/// - `GraphDef`s;
/// - tagged run metadata protos;
/// - summaries with TensorFlow 1.x `image` or `audio`;
/// - summaries with `tensor` set to a rank-1 tensor of type `DT_STRING`;
/// - for audio metadata, summaries with `tensor` set to a shape-`[k, 2]` tensor of type
/// `DT_STRING`, in which case the second axis is assumed to represent string labels and is
/// dropped entirely;
/// - for graph sub-plugin metadata, summaries with `tensor` set to a rank-0 tensor of type
/// `DT_STRING`, which is converted to a shape-`[1]` tensor.
pub fn into_blob_sequence(
self,
metadata: &pb::SummaryMetadata,
Expand Down Expand Up @@ -108,10 +117,7 @@ impl EventValue {
Ok(BlobSequenceValue(tp.string_val))
} else if shape.dim.len() == 2
&& shape.dim[1].size == 2
&& (metadata
.plugin_data
.as_ref()
.map_or(false, |pd| pd.plugin_name == AUDIO_PLUGIN_NAME))
&& is_plugin(&metadata, AUDIO_PLUGIN_NAME)
{
// Extract just the actual audio clips along the first axis.
let audio: Vec<Vec<u8>> = tp
Expand All @@ -120,6 +126,14 @@ impl EventValue {
.map(|chunk| std::mem::take(&mut chunk[0]))
.collect();
Ok(BlobSequenceValue(audio))
} else if shape.dim.is_empty()
&& tp.string_val.len() == 1
&& (is_plugin(&metadata, GRAPH_RUN_METADATA_PLUGIN_NAME)
|| is_plugin(&metadata, GRAPH_RUN_METADATA_WITH_GRAPH_PLUGIN_NAME)
|| is_plugin(&metadata, GRAPH_KERAS_MODEL_PLUGIN_NAME))
{
let data = tp.string_val.into_iter().next().unwrap();
Ok(BlobSequenceValue(vec![data]))
} else {
Err(DataLoss)
}
Expand Down Expand Up @@ -157,6 +171,13 @@ fn tensor_proto_to_scalar(tp: &pb::TensorProto) -> Option<f32> {
}
}

/// Tests whether `md` has plugin name `plugin_name`.
fn is_plugin(md: &pb::SummaryMetadata, plugin_name: &str) -> bool {
md.plugin_data
.as_ref()
.map_or(false, |pd| pd.plugin_name == plugin_name)
}

/// A value from an `Event` whose `graph_def` field is set.
///
/// This contains the raw bytes of a serialized `GraphDef` proto. It implies a fixed tag name and
Expand Down Expand Up @@ -242,7 +263,11 @@ impl SummaryValue {
Some(SCALARS_PLUGIN_NAME) => {
md.data_class = pb::DataClass::Scalar.into();
}
Some(IMAGES_PLUGIN_NAME) | Some(AUDIO_PLUGIN_NAME) => {
Some(IMAGES_PLUGIN_NAME)
| Some(AUDIO_PLUGIN_NAME)
| Some(GRAPH_RUN_METADATA_PLUGIN_NAME)
| Some(GRAPH_RUN_METADATA_WITH_GRAPH_PLUGIN_NAME)
| Some(GRAPH_KERAS_MODEL_PLUGIN_NAME) => {
Comment on lines +268 to +270
Copy link
Contributor

@stephanwlee stephanwlee Jan 19, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Future plan question: are we going to trust the data_class property on the proto in the future? I believe the graph_run* and graph_keras* data are all annotated correctly with the data_class.

Copy link
Contributor Author

@wchargin wchargin Jan 19, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We already trust the data class: the first clause of this match is

        match (md, &*self.0) {
            // Any summary metadata that sets its own data class is expected to already be in the right
            // form.
            (Some(md), _) if md.data_class != i32::from(pb::DataClass::Unknown) => Box::new(md),

(Likewise for Python TensorBoard.)

But I’m not sure if the graph data does set data classes? The data
written by :graphs_demo doesn’t seem to have them, and I don’t see
data classes in the writing code. In any case, data classes were
introduced more recently than the graph_* writing code, so old data
must not have them, I think.

So eventually I would probably like plugins to write summaries with data
class explicitly set, but I expect to maintain this compatibility layer
indefinitely, and that seems okay to me. WDYT?

md.data_class = pb::DataClass::BlobSequence.into();
}
_ => {}
Expand Down Expand Up @@ -671,6 +696,48 @@ mod tests {
);
}

#[test]
fn test_graph_subplugins() {
for &plugin_name in &[
GRAPH_RUN_METADATA_PLUGIN_NAME,
GRAPH_RUN_METADATA_WITH_GRAPH_PLUGIN_NAME,
GRAPH_KERAS_MODEL_PLUGIN_NAME,
] {
let md = pb::SummaryMetadata {
plugin_data: Some(PluginData {
plugin_name: plugin_name.to_string(),
content: b"1".to_vec(),
..Default::default()
}),
..Default::default()
};
let v = SummaryValue(Box::new(Value::Tensor(pb::TensorProto {
dtype: pb::DataType::DtString.into(),
tensor_shape: Some(tensor_shape(&[])),
string_val: vec![b"some-graph-proto".to_vec()],
..Default::default()
})));

// Test both metadata and enrichment here, for convenience.
let initial_metadata = v.initial_metadata(Some(md));
assert_eq!(
*initial_metadata,
pb::SummaryMetadata {
plugin_data: Some(PluginData {
plugin_name: plugin_name.to_string(),
content: b"1".to_vec(),
..Default::default()
}),
data_class: pb::DataClass::BlobSequence.into(),
..Default::default()
},
);
Comment on lines +722 to +734
Copy link
Contributor

@stephanwlee stephanwlee Jan 19, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No AI required: this test feels too mechanical :\

Copy link
Contributor Author

@wchargin wchargin Jan 19, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Like, too much of “assert that f(x) equals the result of manually
computing f(x)”? I agree; it does. I think that I feel okay with it
since it is actually meaningful coverage—dropping a single line in the
match would cause parts of the graph plugin to silently not render.

Or maybe you meant that the test implementation spends a lot of lines on
proto manipulation? Agreed there, too. We could refactor to reuse proto
sub-structures, but I generally like optimizing for readability in tests
like this. In #4552, for example, I made the tests more explicit because
if they had been written like that in the first place, the bug would
likely have been caught in code review: new Date(1610613017) is an
obvious bug—too few digits—whereas new Date(fakePoints1[0].x) is not
obviously correct and just repeats the implementation code, which was
also wrong.

Noted that you said “no AI required”, so feel free to reply or not, as
you like. Glad that this passes your minimum bar, but also always
interested in doing better than that.

let expected_enriched = BlobSequenceValue(vec![b"some-graph-proto".to_vec()]);
let actual_enriched = EventValue::Summary(v).into_blob_sequence(&initial_metadata);
assert_eq!(actual_enriched, Ok(expected_enriched));
}
}

#[test]
fn test_enrich_graph_def() {
let v = EventValue::GraphDef(GraphDefValue(vec![1, 2, 3, 4]));
Expand Down