Skip to content

rust: support audio #4566

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 5 commits into from
Jan 19, 2021
Merged

rust: support audio #4566

Changes from all commits
Commits
Show all changes
5 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
04eaddf
[rust-audio: update diffbase]
wchargin Jan 19, 2021
fe3f131
[rust-audio: 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
156 changes: 148 additions & 8 deletions tensorboard/data/server/data_compat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ use pb::summary_metadata::PluginData;

pub(crate) const SCALARS_PLUGIN_NAME: &str = "scalars";
pub(crate) const IMAGES_PLUGIN_NAME: &str = "images";
pub(crate) const AUDIO_PLUGIN_NAME: &str = "audio";
pub(crate) const GRAPHS_PLUGIN_NAME: &str = "graphs";

/// The inner contents of a single value from an event.
Expand Down Expand Up @@ -71,11 +72,13 @@ impl EventValue {

/// Consumes this event value and enriches it into a blob sequence.
///
/// For now, this supports `GraphDef`s, summaries with `image`, or summaries with `tensor` set
/// to a rank-1 tensor of type `DT_STRING`.
/// For now, this supports `GraphDef`s, 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.
pub fn into_blob_sequence(
self,
Copy link
Contributor

@psybuzz psybuzz Jan 19, 2021

Choose a reason for hiding this comment

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

(no action needed) Initially, I thought it was curious that we don't have to write mut self in the signature, but since the function takes ownership, it is implied that 'self' can be mutated. Is this understanding correct?

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.

Kind of. The mut in x: mut self is a property of the variable
binding, not the type. To demonstrate the orthogonality, consider:

let mut xs = vec![0, 1, 2];

let mut x: &mut i32 = &mut xs[0];
//  ^^^     ~~~

The first mut means that it’s legal to write x = &mut xs[1]:
reassigning to the variable binding. This is entirely up to the call
site, like const/let in JavaScript.

The second mut means that it’s legal to write *x = 0, mutating the
value behind the binding. This is a property of the type, like
ReadonlyArray/Array in TypeScript: if I give you a read-only array,
you can’t just choose to interpret it as an Array.

It can often be useful to have a mutable binding even if the value is at
an immutable reference type:

fn parse_number(mut buf: &str) -> Option<i32> {
    let mut sign = 1;
    if buf.starts_with('-') {
        sign = -1;
        buf = &buf[1..];
        //  ^^^^^^^^^^^  reassigns to `buf`, even though the underlying
        //               buffer's *contents* are immutable
    }
    Some(buf.parse::<i32>().ok()? * sign)
}

The mutability of the binding is always at the discretion of the caller,
and doesn’t affect types or APIs. The caller can change it at will:

let xs = vec![0, 1, 2];
// `xs.push(3)` is illegal here; `xs` is not `mut`

let mut ys = xs; // same value (moved, not copied); binding is now mutable
ys.push(3); // legal

let zs = ys;
// `zs.push(4)` is illegal again

The code in this PR is analogous. We bind self immutably, and we don’t
modify self. But in the Tensor clause of the match expression, we
bind tp mutably, so that we can call chunks_exact_mut.

So your claim that

since the function takes ownership, it is implied that 'self' can be mutated

is correct, in a way. It’s implied that self can be changed to a
mutable binding since the binding is always “owned” by the function.

_metadata: &pb::SummaryMetadata,
metadata: &pb::SummaryMetadata,
) -> Result<BlobSequenceValue, DataLoss> {
match self {
EventValue::GraphDef(GraphDefValue(blob)) => Ok(BlobSequenceValue(vec![blob])),
Expand All @@ -86,11 +89,29 @@ impl EventValue {
let buf = im.encoded_image_string;
Ok(BlobSequenceValue(vec![w, h, buf]))
}
pb::summary::value::Value::Tensor(tp) => {
if tp.dtype == i32::from(pb::DataType::DtString)
&& tp.tensor_shape.map_or(false, |shape| shape.dim.len() == 1)
{
pb::summary::value::Value::Audio(au) => {
Ok(BlobSequenceValue(vec![au.encoded_audio_string]))
}
pb::summary::value::Value::Tensor(mut tp)
if tp.dtype == i32::from(pb::DataType::DtString) =>
{
let shape = tp.tensor_shape.unwrap_or_default();
if shape.dim.len() == 1 {
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))
{
// Extract just the actual audio clips along the first axis.
let audio: Vec<Vec<u8>> = tp
.string_val
.chunks_exact_mut(2)
.map(|chunk| std::mem::take(&mut chunk[0]))
.collect();
Ok(BlobSequenceValue(audio))
} else {
Err(DataLoss)
}
Expand Down Expand Up @@ -189,13 +210,14 @@ impl SummaryValue {
(Some(md), _) if md.data_class != i32::from(pb::DataClass::Unknown) => Box::new(md),
(_, Value::SimpleValue(_)) => blank(SCALARS_PLUGIN_NAME, pb::DataClass::Scalar),
(_, Value::Image(_)) => blank(IMAGES_PLUGIN_NAME, pb::DataClass::BlobSequence),
(_, Value::Audio(_)) => blank(AUDIO_PLUGIN_NAME, pb::DataClass::BlobSequence),
(Some(mut md), _) => {
// Use given metadata, but first set data class based on plugin name, if known.
match md.plugin_data.as_ref().map(|pd| pd.plugin_name.as_str()) {
Some(SCALARS_PLUGIN_NAME) => {
md.data_class = pb::DataClass::Scalar.into();
}
Some(IMAGES_PLUGIN_NAME) => {
Some(IMAGES_PLUGIN_NAME) | Some(AUDIO_PLUGIN_NAME) => {
md.data_class = pb::DataClass::BlobSequence.into();
}
_ => {}
Expand Down Expand Up @@ -552,6 +574,60 @@ mod tests {
);
}

#[test]
fn test_metadata_tf1x_audio() {
let v = SummaryValue(Box::new(Value::Audio(pb::summary::Audio {
sample_rate: 44100.0,
encoded_audio_string: b"RIFFabcd".to_vec(),
..Default::default()
})));
let result = v.initial_metadata(None);

assert_eq!(
*result,
pb::SummaryMetadata {
plugin_data: Some(PluginData {
plugin_name: AUDIO_PLUGIN_NAME.to_string(),
..Default::default()
}),
data_class: pb::DataClass::BlobSequence.into(),
..Default::default()
}
);
}

#[test]
fn test_metadata_tf2x_audio_without_dataclass() {
let md = pb::SummaryMetadata {
plugin_data: Some(PluginData {
plugin_name: AUDIO_PLUGIN_NAME.to_string(),
content: b"preserved!".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(&[1, 2])),
string_val: vec![b"\x89PNGabc".to_vec(), b"label".to_vec()],
..Default::default()
})));
let result = v.initial_metadata(Some(md));

assert_eq!(
*result,
pb::SummaryMetadata {
plugin_data: Some(PluginData {
plugin_name: AUDIO_PLUGIN_NAME.to_string(),
content: b"preserved!".to_vec(),
..Default::default()
}),
data_class: pb::DataClass::BlobSequence.into(),
..Default::default()
}
);
}

#[test]
fn test_enrich_graph_def() {
let v = EventValue::GraphDef(GraphDefValue(vec![1, 2, 3, 4]));
Expand Down Expand Up @@ -672,6 +748,70 @@ mod tests {
Err(DataLoss)
);
}

#[test]
fn test_enrich_tf1x_audio() {
let v = SummaryValue(Box::new(Value::Audio(pb::summary::Audio {
sample_rate: 44100.0,
encoded_audio_string: b"RIFFabcd".to_vec(),
..Default::default()
})));
let md = v.initial_metadata(None);
let expected = BlobSequenceValue(vec![b"RIFFabcd".to_vec()]);
assert_eq!(
EventValue::Summary(v).into_blob_sequence(md.as_ref()),
Ok(expected)
);
}

#[test]
fn test_enrich_audio_without_labels() {
let v = EventValue::Summary(SummaryValue(Box::new(Value::Tensor(pb::TensorProto {
dtype: pb::DataType::DtString.into(),
tensor_shape: Some(tensor_shape(&[3])),
string_val: vec![
b"RIFFwav0".to_vec(),
b"RIFFwav1".to_vec(),
b"RIFFwav2".to_vec(),
],
..Default::default()
}))));
let expected = BlobSequenceValue(vec![
b"RIFFwav0".to_vec(),
b"RIFFwav1".to_vec(),
b"RIFFwav2".to_vec(),
]);
assert_eq!(
v.into_blob_sequence(&blank(AUDIO_PLUGIN_NAME, pb::DataClass::BlobSequence)),
Ok(expected)
);
}

#[test]
fn test_enrich_audio_with_labels() {
let v = EventValue::Summary(SummaryValue(Box::new(Value::Tensor(pb::TensorProto {
dtype: pb::DataType::DtString.into(),
tensor_shape: Some(tensor_shape(&[3, 2])),
string_val: vec![
b"RIFFwav0".to_vec(),
b"label 0".to_vec(),
b"RIFFwav1".to_vec(),
b"label 1".to_vec(),
b"RIFFwav2".to_vec(),
b"label 2".to_vec(),
],
..Default::default()
}))));
let expected = BlobSequenceValue(vec![
b"RIFFwav0".to_vec(),
b"RIFFwav1".to_vec(),
b"RIFFwav2".to_vec(),
]);
assert_eq!(
v.into_blob_sequence(&blank(AUDIO_PLUGIN_NAME, pb::DataClass::BlobSequence)),
Ok(expected)
);
}
}

mod unknown {
Expand Down