Commit edea2b67 authored by Terry Koo's avatar Terry Koo
Browse files

Remove runtime because reasons.

parent a4bb31d0
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#ifndef DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_
#define DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_
#include <map>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "dragnn/protos/export.pb.h"
#include "dragnn/runtime/trained_model.h"
#include "syntaxnet/base.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// Converter that extracts the cell computation from a DRAGNN component and
// writes it as a Myelin Flow.
//
// The trained model that contains the DRAGNN component must also contain a
// CellSubgraphSpec proto embedded into the TF graph as a specifically-named
// constant node (see runtime_support.py). The CellSubgraphSpec defines the
// boundaries of the cell comptation.
//
// The converted Myelin Flow contains a single function that runs the cell and
// is named after the component. The function inputs and outputs are reference
// variables, so they can be pointed at externally-managed pieces of memory,
// provided sufficient size and alignment. The function inputs and outputs are
// marked with special aliases, namely:
// INPUT/<CellSubgraphSpec.Input.name>
// OUTPUT/<CellSubgraphSpec.Output.name>
class MyelinCellConverter {
public:
// Extracts the cell of the DRAGNN component named |component_name| from the
// |trained_model| and overwrites the |flow| with an equivalent Myelin Flow.
// The |flow| file output is deterministic given identical inputs. On error,
// returns non-OK.
static tensorflow::Status Convert(const string &component_name,
const TrainedModel &trained_model,
string *flow);
private:
// A (node_name, output_index) pair denoting a tensor.
using TensorId = std::pair<string, uint32>;
// Flow file writer; defined in the .cc file.
class Writer;
// An operation that makes up the cell, convertible to a Myelin operation.
struct Operation {
// The TF graph node represented by this operation.
const tensorflow::NodeDef *node = nullptr;
// Myelin variable names of inputs to this operation. Order matters.
std::vector<string> inputs;
// Number of outputs observed for this operation. Some of the outputs in
// [0,|num_outputs|) might not actually be used in the cell, but we must
// create variables for all of them to match the expected output arity and
// ordering of the operation.
uint32 num_outputs = 0;
};
// Creates an empty converter.
MyelinCellConverter() = default;
// Implements the static Convert() method.
tensorflow::Status ConvertImpl(const string &component_name,
const TrainedModel &trained_model,
string *flow);
// Populates the |inputs_| and |outputs_| based on the |spec|. On error,
// returns non-OK.
tensorflow::Status BuildInputsAndOutputs(const CellSubgraphSpec &spec);
// Walks from the |outputs_| to the |inputs_| in the |trained_model_|, adding
// to |operations_| along the way. Requires that BuildInputsAndOutputs() was
// called. On error, returns non-OK.
tensorflow::Status BuildOperations();
// Writes each section of a flow file to the |writer|.
tensorflow::Status WriteVariables(Writer *writer) const;
void WriteOperations(Writer *writer) const;
void WriteFunctions(Writer *writer) const;
void WriteConnectors(Writer *writer) const;
void WriteBlobs(Writer *writer) const;
// Writes a variable for the |output_index|'th output of the |node| to the
// |writer|. Retrieves constant variable data from the |trained_model_| if
// necessary. On error, returns non-OK.
tensorflow::Status WriteVariable(const tensorflow::NodeDef &node,
uint32 output_index, Writer *writer) const;
// Writes the |operation| to the |writer|.
void WriteOperation(const Operation &operation, Writer *writer) const;
// Returns the set of aliases associated with the |tensor_id|.
std::set<string> GetAliases(const TensorId &tensor_id) const;
// Parses a |tensor_name| into a |tensor_id|. E.g.,
// "foo/bar:1" => ("foo/bar", 1)
// "baz" => ("baz", 0)
// On error, returns non-OK. It is an error if the |tensor_name| denotes a
// control dependency.
static tensorflow::Status ParseTensorId(const string &tensor_name,
TensorId *tensor_id);
// Returns the canonically-formatted name of the Myelin variable associated
// with the |tensor_id|.
static string AsVariableName(const TensorId &tensor_id);
// Name of the component being converted.
string component_name_;
// Trained model that contains the DRAGNN model.
const TrainedModel *trained_model_ = nullptr;
// Mapping from input tensor to logical input name.
std::map<TensorId, string> inputs_;
// Mapping from output tensor to logical output names. There may be more than
// one name due to layer aliases (e.g., "last_layer").
std::map<TensorId, std::set<string>> outputs_;
// Mapping from node name to Operation.
std::map<string, Operation> operations_;
};
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELIN_CELL_CONVERTER_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_cell_converter.h"
#include <string.h>
#include <iostream>
#include <string>
#include <utility>
#include <vector>
#include "dragnn/core/test/generic.h"
#include "dragnn/runtime/alignment.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/trained_model.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "sling/myelin/flow.h"
#include "sling/myelin/graph.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Relative path to a saved model.
constexpr char kSavedModelDir[] = "dragnn/runtime/testdata/rnn_tagger";
// Names of components in the saved model.
const char *kComponentNames[] = {"rnn", "tagger"};
// Returns a valid saved model directory.
string GetSavedModelDir() {
return tensorflow::io::JoinPath(test::GetTestDataPrefix(), kSavedModelDir);
}
// Returns a string like "1048576 bytes (1.0MiB)".
string FormatSize(int64 size) {
return tensorflow::strings::StrCat(
size, " bytes (", tensorflow::strings::HumanReadableNumBytes(size), ")");
}
// Logs the |flow|, using the |description| in the log messages.
void DumpFlow(const sling::myelin::Flow &flow, const string &description) {
VLOG(1) << description << " Flow:\n" << flow.ToString();
// Log messages are truncated when they get too long. Dump the DOT file to
// stdout so we get the whole thing.
std::cout << description << " DOT:\n"
<< sling::myelin::FlowToDotGraph(flow, {}) << std::endl;
}
// Returns true if the |variable| is a function input or output.
bool IsFunctionInputOrOutput(const sling::myelin::Flow::Variable &variable) {
// Inputs and outputs are marked with special aliases.
for (tensorflow::StringPiece alias : variable.aliases) {
if (tensorflow::str_util::StartsWith(alias, "INPUT/")) return true;
if (tensorflow::str_util::StartsWith(alias, "OUTPUT/")) return true;
}
return false;
}
// Returns a list of (tensor,array) pairs, one for each input and output of the
// |flow| and |network|. The arrays are zero-filled.
std::vector<std::pair<sling::myelin::Tensor *, UniqueAlignedArray>>
GetInputAndOutputTensorsAndArrays(const sling::myelin::Flow &flow,
const sling::myelin::Network &network) {
std::vector<std::pair<sling::myelin::Tensor *, UniqueAlignedArray>>
tensors_and_arrays;
for (const sling::myelin::Flow::Variable *variable : flow.vars()) {
// NB: Gating on variable->in || variable->out is too coarse, because that
// also includes constants.
if (!IsFunctionInputOrOutput(*variable)) continue;
sling::myelin::Tensor *tensor = network.GetParameter(variable->name);
CHECK(tensor != nullptr)
<< "Failed to find tensor for variable " << variable->name;
// Currently, inputs and outputs are either int32 or float, which are the
// same size and have the same representation of zero. Therefore, we can
// treat them the same at the byte level.
CHECK(tensor->type() == sling::myelin::DT_FLOAT ||
tensor->type() == sling::myelin::DT_INT32);
static_assert(sizeof(int32) == sizeof(float), "Unexpected size mismatch");
const int bytes = variable->shape.elements() * sizeof(float);
UniqueAlignedArray array;
array.Reset(bytes);
memset(array.view().data(), 0, bytes); // zero for int32 or float
tensors_and_arrays.emplace_back(tensor, std::move(array));
VLOG(1) << "Created array of " << bytes << " bytes for variable "
<< variable->name << " with aliases "
<< tensorflow::str_util::Join(variable->aliases, ", ");
}
return tensors_and_arrays;
}
// Loads a trained model, converts it into a Flow, and then analyzes, compiles,
// and runs the Flow.
TEST(MyelinCellConverterTest, LoadConvertAndRun) {
TrainedModel trained_model;
TF_ASSERT_OK(trained_model.Reset(GetSavedModelDir()));
for (const string component_name : kComponentNames) {
LOG(INFO) << "Component: " << component_name;
string data;
TF_ASSERT_OK(
MyelinCellConverter::Convert(component_name, trained_model, &data));
LOG(INFO) << component_name << " flow size = " << FormatSize(data.size());
sling::myelin::Flow flow;
flow.Read(data.data(), data.size());
sling::myelin::Library library;
RegisterMyelinLibraries(&library);
DumpFlow(flow, tensorflow::strings::StrCat(component_name, " original"));
flow.Analyze(library);
DumpFlow(flow, tensorflow::strings::StrCat(component_name, " analyzed"));
sling::myelin::Network network;
ASSERT_TRUE(network.Compile(flow, library));
const sling::myelin::Cell *cell = network.GetCell(component_name);
ASSERT_TRUE(cell != nullptr);
LOG(INFO) << component_name
<< " code size = " << FormatSize(cell->code().size());
// Create an instance and point input/output references at arrays.
sling::myelin::Instance instance(cell);
const auto tensors_and_arrays =
GetInputAndOutputTensorsAndArrays(flow, network);
for (const auto &tensor_and_array : tensors_and_arrays) {
instance.SetReference(tensor_and_array.first,
tensor_and_array.second.view().data());
}
// This is just a "don't crash" test. Myelin behavior will be exercised
// more thoroughly in regression tests.
LOG(INFO) << "Running " << component_name;
instance.Compute();
LOG(INFO) << "Successfully ran " << component_name << "!";
}
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include <stddef.h>
#include <string>
#include "dragnn/core/compute_session.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/protos/trace.pb.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/fixed_embeddings.h"
#include "dragnn/runtime/linked_embeddings.h"
#include "dragnn/runtime/math/types.h"
#include "dragnn/runtime/myelin/myelin_dynamic_component_base.h"
#include "dragnn/runtime/network_states.h"
#include "dragnn/runtime/session_state.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// A Myelin-based version of DynamicComponent.
//
// This implementation of MyelinDynamicComponentBase has the most generality
// w.r.t. input features and links, but suffers from ComputeSession overhead.
class MyelinDynamicComponent : public MyelinDynamicComponentBase {
public:
// Implements Component.
tensorflow::Status Evaluate(SessionState *session_state,
ComputeSession *compute_session,
ComponentTrace *component_trace) const override;
protected:
// Unlike other specializations, this component will only be active if the
// spec is explicitly modified to support Myelin (and flow resources are
// generated).
bool Supports(const ComponentSpec &spec,
const string &normalized_builder_name) const override {
return normalized_builder_name == "MyelinDynamicComponent";
}
bool PreferredTo(const Component &other) const override { return false; }
private:
// Forbid batches and beams.
static constexpr int kEvaluateNumItems = 1;
};
tensorflow::Status MyelinDynamicComponent::Evaluate(
SessionState *session_state, ComputeSession *compute_session,
ComponentTrace *component_trace) const {
NetworkStates &network_states = session_state->network_states;
FixedEmbeddings &fixed_embeddings = GetFixedEmbeddings(session_state);
LinkedEmbeddings &linked_embeddings = GetLinkedEmbeddings(session_state);
sling::myelin::Instance &instance = GetInstance(session_state);
for (size_t step_index = 0; !compute_session->IsTerminal(name());
++step_index) {
network_states.AddStep();
TF_RETURN_IF_ERROR(fixed_embeddings.Reset(&fixed_embedding_manager(),
network_states, compute_session));
TF_RETURN_IF_ERROR(linked_embeddings.Reset(
&linked_embedding_manager(), network_states, compute_session));
// Bind inputs and outputs into the |instance|.
BindInputIds(fixed_embeddings, &instance);
BindInputLinks(linked_embeddings, &instance);
BindInputRecurrences(step_index, network_states, &instance);
BindOutputLayers(step_index, network_states, &instance);
// Invoke the cell in the |instance|.
instance.Compute();
MaybeTrace(step_index, &instance, component_trace);
// If the component is deterministic, take the oracle transition instead of
// predicting the next transition using the logits.
if (deterministic()) {
compute_session->AdvanceFromOracle(name());
} else {
// AddStep() may invalidate the logits (due to reallocation), so the layer
// lookup cannot be hoisted out of this loop.
const Vector<float> logits(
network_states.GetLayer(logits_handle()).row(step_index));
if (!compute_session->AdvanceFromPrediction(
name(), logits.data(), kEvaluateNumItems, logits.size())) {
return tensorflow::errors::Internal(
"Error in ComputeSession::AdvanceFromPrediction()");
}
}
}
return tensorflow::Status::OK();
}
DRAGNN_RUNTIME_REGISTER_COMPONENT(MyelinDynamicComponent);
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_dynamic_component_base.h"
#include <string.h>
#include <algorithm>
#include <set>
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/transition_system_traits.h"
#include "tensorflow/core/lib/core/errors.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
constexpr char MyelinDynamicComponentBase::kLogitsName[];
tensorflow::Status MyelinDynamicComponentBase::Validate(
const ComponentSpec &component_spec) {
if (!component_spec.attention_component().empty()) {
return tensorflow::errors::Unimplemented("Attention is not supported");
}
for (const auto &fixed_feature : component_spec.fixed_feature()) {
if (fixed_feature.embedding_dim() != -1) {
return tensorflow::errors::InvalidArgument(
"Myelin requires non-embedded fixed features");
}
}
for (const auto &linked_feature : component_spec.linked_feature()) {
if (linked_feature.embedding_dim() != -1) {
return tensorflow::errors::InvalidArgument(
"Myelin requires non-multiplied linked features");
}
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::LookupVector(
const string &name, sling::myelin::Type type, int dimension,
sling::myelin::Tensor **vector) const {
*vector = nullptr; // so it is null if we error out
sling::myelin::Tensor *tensor = network_.GetParameter(name);
if (tensor == nullptr) {
return tensorflow::errors::NotFound("No Myelin tensor named '", name, "'");
}
if (tensor->type() != type) {
return tensorflow::errors::InvalidArgument(
"Myelin tensor has wrong type '", name, "' ", tensor->TypeString(),
" (expected ", sling::myelin::TypeTraits::of(type).name(), ")");
}
int num_nontrivial_dims = 0;
for (int i = 0; i < tensor->rank(); ++i) {
if (tensor->dim(i) > 1) ++num_nontrivial_dims;
}
if (num_nontrivial_dims > 1) {
return tensorflow::errors::InvalidArgument(
"Myelin tensor has non-vector-like shape: '", name, "' ",
tensor->TypeString());
}
// Since the |tensor| is vector-like, elements() is equivalent to the vector
// dimension and smooths over edges like rank=0.
if (dimension >= 0 && tensor->elements() != dimension) {
return tensorflow::errors::InvalidArgument(
"Myelin vector has the wrong dimension '", name, "' ",
tensor->TypeString(), " (expected ", dimension, ")");
}
if (internal::kAlignmentBytes % tensor->byte_alignment() != 0) {
return tensorflow::errors::FailedPrecondition(
"Myelin vector '", name, "' has incompatible byte alignment ",
tensor->byte_alignment(), " (vs ", internal::kAlignmentBytes, ")");
}
for (int i = 0; i < tensor->rank(); ++i) {
if (internal::kAlignmentBytes % tensor->minalign(i) != 0) {
return tensorflow::errors::FailedPrecondition(
"Myelin vector '", name, "' has incompatible minimum alignment ",
tensor->minalign(i), " for dimension ", i, " (vs ",
internal::kAlignmentBytes, ")");
}
}
// Success; update |vector| to non-null.
*vector = tensor;
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::InitializeInputIds() {
const int num_channels = fixed_embedding_manager_.num_channels();
input_ids_.resize(fixed_embedding_manager_.num_embeddings());
for (int channel_id = 0; channel_id < num_channels; ++channel_id) {
DCHECK(!fixed_embedding_manager_.is_embedded(channel_id));
const int channel_base = fixed_embedding_manager_.channel_base(channel_id);
const int channel_size = fixed_embedding_manager_.channel_size(channel_id);
for (int index = 0; index < channel_size; ++index) {
InputId &input = input_ids_[channel_base + index];
const string name = MakeMyelinInputFixedFeatureIdName(channel_id, index);
TF_RETURN_IF_ERROR(
LookupVector(name, sling::myelin::DT_INT32, 1, &input.id));
VLOG(1) << "Component '" << name_ << "' fixed channel " << channel_id
<< " index " << index << ": Added feature ID";
}
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::InitializeInputLinks() {
const int num_channels = linked_embedding_manager_.num_channels();
input_links_.resize(num_channels);
for (int channel_id = 0; channel_id < num_channels; ++channel_id) {
InputLink &input = input_links_[channel_id];
const int dimension = linked_embedding_manager_.embedding_dim(channel_id);
const string activations_name =
MakeMyelinInputLinkedActivationVectorName(channel_id);
const string out_of_bounds_name =
MakeMyelinInputLinkedOutOfBoundsIndicatorName(channel_id);
TF_RETURN_IF_ERROR(LookupVector(activations_name, sling::myelin::DT_FLOAT,
dimension, &input.activations));
VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
<< ": Added activations";
// Allow NOT_FOUND, for linked embedding channels that don't multiply the
// input activations with an embedding matrix.
const tensorflow::Status status = LookupVector(
out_of_bounds_name, sling::myelin::DT_FLOAT, 1, &input.out_of_bounds);
if (status.ok()) {
VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
<< ": Added out-of-bounds indicator for multiplication";
} else if (status.code() == tensorflow::error::NOT_FOUND) {
VLOG(1) << "Component '" << name_ << "' linked channel " << channel_id
<< ": No out-of-bounds indicator; not multiplied";
} else {
return status;
}
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::InitializeInputRecurrences(
const sling::myelin::Flow &flow, const NetworkStateManager &manager) {
for (const string &layer_name : GetRecurrentLayerNames(flow)) {
input_recurrences_.emplace_back();
InputRecurrence &input = input_recurrences_.back();
const string name = MakeMyelinInputRecurrentLayerName(layer_name);
size_t dimension = 1;
TF_RETURN_IF_ERROR(
manager.LookupLayer(name_, layer_name, &dimension, &input.handle));
TF_RETURN_IF_ERROR(LookupVector(name, sling::myelin::DT_FLOAT, dimension,
&input.previous_output));
VLOG(1) << "Component '" << name_ << "' recurrence '" << layer_name
<< "': Added link to previous output";
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::InitializeOutputLayers(
const sling::myelin::Flow &flow, NetworkStateManager *manager) {
// Mapping from cell output tensor to layer name, for detecting layer aliases.
std::map<const sling::myelin::Tensor *, string> tensor_to_layer;
for (const string &layer_name : GetOutputLayerNames(flow)) {
output_layers_.emplace_back();
OutputLayer &output = output_layers_.back();
const string name = MakeMyelinOutputLayerName(layer_name);
TF_RETURN_IF_ERROR(
LookupVector(name, sling::myelin::DT_FLOAT, -1, &output.layer));
// Add a new output layer or create an alias to an existing one.
if (tensor_to_layer.find(output.layer) == tensor_to_layer.end()) {
tensor_to_layer[output.layer] = layer_name;
const size_t dimension = output.layer->elements();
TF_RETURN_IF_ERROR(
manager->AddLayer(layer_name, dimension, &output.handle));
VLOG(1) << "Component '" << name_ << "' output '" << layer_name
<< "': Added new layer";
} else {
const string &original_name = tensor_to_layer[output.layer];
output_layers_.pop_back(); // not a "real" output
TF_RETURN_IF_ERROR(manager->AddLayerAlias(layer_name, original_name));
VLOG(1) << "Component '" << name_ << "' output '" << layer_name
<< "': Alias of '" << original_name << "'";
}
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::InitializeConstantVectors() {
// Find the maximum recurrent layer dimension; the |zeros_| must be this big.
int max_dimension = 1; // ensure at least one element, for |zero_|
for (const InputRecurrence &input : input_recurrences_) {
max_dimension = std::max(max_dimension, input.previous_output->elements());
}
// Allocate the backing array and parcel it out into sub-views.
const std::vector<size_t> sizes = {sizeof(float),
max_dimension * sizeof(float)};
array_.Reset(ComputeTotalBytesWithAlignmentPadding(sizes));
memset(array_.view().data(), 0, array_.view().size()); // = 0.0 for float
std::vector<MutableAlignedView> views;
TF_RETURN_IF_ERROR(array_.view().Split(sizes, &views));
DCHECK_EQ(views.size(), 2);
// Promote to typed vectors.
one_ = Vector<float>(views[0]);
zero_ = Vector<float>(views[1], 1);
zeros_ = Vector<float>(views[1]);
DCHECK_EQ(zero_.size(), 1);
DCHECK_EQ(one_.size(), 1);
DCHECK_EQ(zeros_.size(), max_dimension);
// All memory was already zeroed, so only |one_| needs to be initialized.
MutableVector<float> mutable_one(views[0]);
mutable_one[0] = 1.0;
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::MaybeInitializeLogits(
const ComponentSpec &component_spec, const NetworkStateManager &manager) {
// Logits are unnecessary when the component is deterministic.
deterministic_ = TransitionSystemTraits(component_spec).is_deterministic;
if (deterministic_) return tensorflow::Status::OK();
size_t dimension = 0;
TF_RETURN_IF_ERROR(
manager.LookupLayer(name_, kLogitsName, &dimension, &logits_handle_));
if (dimension != component_spec.num_actions()) {
return tensorflow::errors::InvalidArgument(
"Dimension mismatch between classification logits (", dimension,
") and ComponentSpec.num_actions (", component_spec.num_actions(),
") in component '", name_, "'");
}
return tensorflow::Status::OK();
}
tensorflow::Status MyelinDynamicComponentBase::Initialize(
const ComponentSpec &component_spec, VariableStore *variable_store,
NetworkStateManager *network_state_manager,
ExtensionManager *extension_manager) {
name_ = component_spec.name();
TF_RETURN_IF_ERROR(Validate(component_spec));
const Resource *resource = nullptr;
TF_RETURN_IF_ERROR(LookupMyelinFlowResource(component_spec, &resource));
const string &flow_path = resource->part(0).file_pattern();
sling::myelin::Flow flow;
TF_RETURN_IF_ERROR(LoadMyelinFlow(flow_path, &flow));
VLOG(1) << "Original Flow for '" << name_ << "':\n" << flow.ToString();
// TODO(googleuser): Add support for optional profiling, via something like:
// if (...) network_.set_profiling(true)
// network_.Compile(flow, library);
// ...
// instance->Compute();
// sling::myelin::Profile profile(instance);
// VLOG(1) << profile.ASCIIReport();
RegisterMyelinLibraries(&library_);
flow.Analyze(library_);
VLOG(1) << "Analyzed Flow for '" << name_ << "':\n" << flow.ToString();
if (!network_.Compile(flow, library_)) {
return tensorflow::errors::Internal(
"Failed to compile Myelin network for component '", name_, "'");
}
cell_ = network_.GetCell(name_);
if (cell_ == nullptr) {
return tensorflow::errors::FailedPrecondition(
"No function named '", name_, "' in Myelin network for component '",
name_, "'");
}
VLOG(1) << name_ << ": " << cell_->code().size() << " bytes of Myelin code";
// Configure the inputs and outputs of the Myelin cell. As with NetworkUnit
// and NetworkUnitBase, output layers and input features must be initialized
// in a particular order to enable recurrent inputs. Specifically, we must
// populate output layers first, so they are available for recurrent access,
// both by the |input_recurrences_| and the |linked_embedding_manager_|.
TF_RETURN_IF_ERROR(InitializeOutputLayers(flow, network_state_manager));
TF_RETURN_IF_ERROR(fixed_embedding_manager_.Reset(
component_spec, variable_store, network_state_manager));
TF_RETURN_IF_ERROR(linked_embedding_manager_.Reset(
component_spec, variable_store, network_state_manager));
TF_RETURN_IF_ERROR(InitializeInputIds());
TF_RETURN_IF_ERROR(InitializeInputLinks());
TF_RETURN_IF_ERROR(InitializeInputRecurrences(flow, *network_state_manager));
TF_RETURN_IF_ERROR(InitializeConstantVectors());
TF_RETURN_IF_ERROR(
MaybeInitializeLogits(component_spec, *network_state_manager));
extension_manager->GetShared(&fixed_embeddings_handle_);
extension_manager->GetShared(&linked_embeddings_handle_);
extension_manager->AddLocal(&instance_handle_);
return tensorflow::Status::OK();
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#ifndef DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_
#define DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_
#include <stddef.h>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "dragnn/protos/cell_trace.pb.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/protos/trace.pb.h"
#include "dragnn/runtime/alignment.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/extensions.h"
#include "dragnn/runtime/fixed_embeddings.h"
#include "dragnn/runtime/linked_embeddings.h"
#include "dragnn/runtime/math/types.h"
#include "dragnn/runtime/myelin/myelin_tracing.h"
#include "dragnn/runtime/network_states.h"
#include "dragnn/runtime/session_state.h"
#include "dragnn/runtime/variable_store.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "sling/myelin/flow.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/dynamic_annotations.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// Base class for Myelin-based versions of DynamicComponent.
//
// Roughly, this is a base class for a version of DynamicComponent where the
// per-transition-step computation is performed by a Myelin cell instead of a
// NetworkUnit. This class implements Initialize() and provides methods that
// can be useful for inference, but does not implement Evaluate().
//
// At initialization time, this class creates lists of configuration structs
// that associate each input or output of the Myelin cell with an operand that
// the DRAGNN runtime manages. See, e.g., InputId and InitializeInputIds().
//
// At inference time, subclasses can bind the relevant DRAGNN runtime operands
// to the inputs and outputs of the Myelin instance (see, e.g., BindInputIds())
// and evaluate the Myelin cell. Like DynamicComponent, the cell should be
// evaluated once per transition and the results used to advance the transition
// system state.
//
// Except as noted below, this is a drop-in replacement for DynamicComponent:
// * The name of the logits layer is hard-coded (see kLogitsName).
// * The fixed and linked channels must have embedding_dim=-1, because the fixed
// lookups and linked multiplications are handled within Myelin.
//
// The MyelinDynamicComponent subclass provides a general-purpose implementation
// of Evaluate(). Other subclasses provide optimized implementations subject to
// restrictions on the possible network configuration.
class MyelinDynamicComponentBase : public Component {
public:
// Partially implements Component.
tensorflow::Status Initialize(const ComponentSpec &component_spec,
VariableStore *variable_store,
NetworkStateManager *network_state_manager,
ExtensionManager *extension_manager) override;
protected:
// Configuration for a fixed feature ID input.
//
// TODO(googleuser): Consider making singleton inputs like the feature ID and
// out-of-bounds indicator into plain value inputs instead of references; it
// is equally fast to copy the value.
struct InputId {
// Tensor to feed with the fixed feature ID.
sling::myelin::Tensor *id = nullptr;
};
// Configuration for a linked feature embedding input.
struct InputLink {
// Tensor to feed with the linked activation vector.
sling::myelin::Tensor *activations = nullptr;
// Tensor to feed with the linked out-of-bounds indicator, or null if the
// embedding does not need to be multiplied.
sling::myelin::Tensor *out_of_bounds = nullptr;
};
// Configuration for a recurrent input.
struct InputRecurrence {
// Handle of the output layer that is recurrently fed back.
LayerHandle<float> handle;
// Tensor to feed with the previous output activation vector.
sling::myelin::Tensor *previous_output = nullptr;
};
// Configuration for an output layer.
struct OutputLayer {
// Handle of the output layer.
LayerHandle<float> handle;
// Tensor that writes to the layer.
sling::myelin::Tensor *layer = nullptr;
};
// Name of the layer containing logits. Unlike DynamicComponent, this class
// does not use the NetworkUnit abstraction and assumes that the logits will
// be stored in this layer.
// TODO(googleuser): Make this configurable, if needed. The logits layer could
// be given a special alias, for example.
static constexpr char kLogitsName[] = "logits";
// Points the cell input |tensor| in the |instance| at the |vector|.
template <class T>
static void BindInput(Vector<T> vector, sling::myelin::Tensor *tensor,
sling::myelin::Instance *instance);
// Points the cell output |tensor| in the |instance| at the |vector|.
template <class T>
static void BindOutput(MutableVector<T> vector, sling::myelin::Tensor *tensor,
sling::myelin::Instance *instance);
// Binds the feature IDs in the |fixed_embeddings| to the |instance| as
// configured by the |input_ids_|.
void BindInputIds(const FixedEmbeddings &fixed_embeddings,
sling::myelin::Instance *instance) const;
// Binds the |embedding| and, if applicable, |is_out_of_bounds| to the
// |input_link| in the |instance|.
void BindInputLink(Vector<float> embedding, bool is_out_of_bounds,
const InputLink &input_link,
sling::myelin::Instance *instance) const;
// Binds the activation vectors in the |linked_embeddings| to the |instance|
// as configured by the |input_links_|.
void BindInputLinks(const LinkedEmbeddings &linked_embeddings,
sling::myelin::Instance *instance) const;
// Binds the output of the step before |step_index| in the |network_states| to
// the |instance| as configured by the |input_recurrences_|.
void BindInputRecurrences(size_t step_index,
const NetworkStates &network_states,
sling::myelin::Instance *instance) const;
// Binds the output layers for the |step_index| in the |network_states| to the
// |instance| as configured by the |output_layers_|.
void BindOutputLayers(size_t step_index, const NetworkStates &network_states,
sling::myelin::Instance *instance) const;
// Returns the reusable fixed and linked embeddings in the |session_state|.
FixedEmbeddings &GetFixedEmbeddings(SessionState *session_state) const;
LinkedEmbeddings &GetLinkedEmbeddings(SessionState *session_state) const;
// Returns the reusable Myelin instance in the |session_state|.
sling::myelin::Instance &GetInstance(SessionState *session_state) const;
// If |component_trace| is non-null, ensures that |step_index|+1 steps exist
// and traces the |instance| in the |step_index|'th step.
void MaybeTrace(size_t step_index, sling::myelin::Instance *instance,
ComponentTrace *component_trace) const;
// Accessors.
const string &name() const { return name_; }
const FixedEmbeddingManager &fixed_embedding_manager() const {
return fixed_embedding_manager_;
}
const LinkedEmbeddingManager &linked_embedding_manager() const {
return linked_embedding_manager_;
}
const sling::myelin::Cell *cell() const { return cell_; }
const std::vector<InputId> &input_ids() const { return input_ids_; }
const std::vector<InputLink> &input_links() const { return input_links_; }
const std::vector<InputRecurrence> &input_recurrences() const {
return input_recurrences_;
}
const std::vector<OutputLayer> &output_layers() const {
return output_layers_;
}
bool deterministic() const { return deterministic_; }
LayerHandle<float> logits_handle() const { return logits_handle_; }
private:
// Returns non-OK if the |component_spec| specifies any unsupported settings.
// This includes both settings that are not yet implemented and those that are
// fundamentally incompatible with this class.
static tensorflow::Status Validate(const ComponentSpec &component_spec);
// Points the |vector| at the variable in the |network_| named |name|, which
// must have a vector-like shape (i.e., having at most one dimension > 1) and
// must match the |type|. If the |dimension| is >= 0, then the |vector| must
// be the same size. On error, returns non-OK and sets |vector| to nullptr.
// Returns NOT_FOUND iff the |name| does not name a variable.
tensorflow::Status LookupVector(const string &name, sling::myelin::Type type,
int dimension,
sling::myelin::Tensor **vector) const;
// Initializes the |input_ids_| based on the |fixed_embedding_manager_| and
// |network_|. On error, returns non-OK.
tensorflow::Status InitializeInputIds();
// Initializes the |input_links_| based on the |linked_embedding_manager_| and
// |network_|. On error, returns non-OK.
tensorflow::Status InitializeInputLinks();
// Initializes the |input_recurrences_| based on the |flow|, |manager|, and
// |network_|. Requires that layers have been added to the |manager|. On
// error, returns non-OK.
tensorflow::Status InitializeInputRecurrences(
const sling::myelin::Flow &flow, const NetworkStateManager &manager);
// Initializes the |output_layers_| based on the |flow|, |manager|, and
// |network_|. Adds layers to the |manager|. On error, returns non-OK.
tensorflow::Status InitializeOutputLayers(const sling::myelin::Flow &flow,
NetworkStateManager *manager);
// Initializes the constant vectors (|zero_|, |one_|, and |zeros_|) and their
// backing |array_|. Requires that the |input_recurrences_| are initialized.
tensorflow::Status InitializeConstantVectors();
// Initializes the |logits_handle_| based on the |component_spec| and
// |manager|, if needed.
tensorflow::Status MaybeInitializeLogits(const ComponentSpec &component_spec,
const NetworkStateManager &manager);
// Name of this component.
string name_;
// Managers for the fixed and linked embeddings used by the component.
FixedEmbeddingManager fixed_embedding_manager_;
LinkedEmbeddingManager linked_embedding_manager_;
// Fixed and linked embeddings.
SharedExtensionHandle<FixedEmbeddings> fixed_embeddings_handle_;
SharedExtensionHandle<LinkedEmbeddings> linked_embeddings_handle_;
// Library of Myelin kernels and transformations.
sling::myelin::Library library_;
// Myelin network that implements the cell computation.
sling::myelin::Network network_;
// Cell that contains the compiled code for this component.
const sling::myelin::Cell *cell_ = nullptr;
// List of fixed feature ID inputs, aligned with the relevant FixedEmbeddings.
std::vector<InputId> input_ids_;
// List of linked feature inputs, aligned with the relevant LinkedEmbeddings.
std::vector<InputLink> input_links_;
// List of recurrent input, not ordered.
std::vector<InputRecurrence> input_recurrences_;
// List of output layers, not ordered.
std::vector<OutputLayer> output_layers_;
// A few constant vectors and their backing array.
UniqueAlignedArray array_;
Vector<float> zero_; // [0.0], for linked out-of-bounds indicators
Vector<float> one_; // [1.0], for linked out-of-bounds indicators
Vector<float> zeros_; // [0.0...0.0], for linked activation vectors
// Whether the transition system is deterministic.
bool deterministic_ = false;
// Handle to the classification logits. Valid iff |deterministic_| is false.
LayerHandle<float> logits_handle_;
// Instance used to evaluate the network cell. Local, since each component
// can have a different cell.
LocalExtensionHandle<sling::myelin::Instance> instance_handle_;
};
// Implementation details below.
template <class T>
void MyelinDynamicComponentBase::BindInput(Vector<T> vector,
sling::myelin::Tensor *tensor,
sling::myelin::Instance *instance) {
// Since Myelin only consumes non-const pointers, const_cast() is required.
// Myelin will not modify the contents of the |vector|, provided it is bound
// to a cell input.
DCHECK(tensor->in()) << tensor->name();
DCHECK(!tensor->out()) << tensor->name();
DCHECK_LE(tensor->elements(), vector.size()) << tensor->name();
instance->SetReference(
tensor,
const_cast<char *>(reinterpret_cast<const char *>(vector.data())));
}
template <class T>
void MyelinDynamicComponentBase::BindOutput(MutableVector<T> vector,
sling::myelin::Tensor *tensor,
sling::myelin::Instance *instance) {
DCHECK(tensor->out()) << tensor->name();
DCHECK_EQ(tensor->elements(), vector.size()) << tensor->name();
instance->SetReference(tensor, reinterpret_cast<char *>(vector.data()));
TF_ANNOTATE_MEMORY_IS_INITIALIZED(vector.data(), vector.size() * sizeof(T));
}
inline void MyelinDynamicComponentBase::BindInputIds(
const FixedEmbeddings &fixed_embeddings,
sling::myelin::Instance *instance) const {
for (size_t i = 0; i < input_ids_.size(); ++i) {
BindInput(fixed_embeddings.ids(i), input_ids_[i].id, instance);
}
}
inline void MyelinDynamicComponentBase::BindInputLink(
Vector<float> embedding, bool is_out_of_bounds, const InputLink &input_link,
sling::myelin::Instance *instance) const {
BindInput(embedding, input_link.activations, instance);
if (input_link.out_of_bounds != nullptr) {
BindInput(is_out_of_bounds ? one_ : zero_, input_link.out_of_bounds,
instance);
}
}
inline void MyelinDynamicComponentBase::BindInputLinks(
const LinkedEmbeddings &linked_embeddings,
sling::myelin::Instance *instance) const {
for (size_t channel_id = 0; channel_id < input_links_.size(); ++channel_id) {
BindInputLink(linked_embeddings.embedding(channel_id),
linked_embeddings.is_out_of_bounds(channel_id),
input_links_[channel_id], instance);
}
}
inline void MyelinDynamicComponentBase::BindInputRecurrences(
size_t step_index, const NetworkStates &network_states,
sling::myelin::Instance *instance) const {
for (const InputRecurrence &input : input_recurrences_) {
if (step_index == 0) {
// The previous output is out-of-bounds, so feed a zero vector. Recall
// that |zeros_| was constructed to be large enough for any recurrence.
BindInput(zeros_, input.previous_output, instance);
} else {
BindInput(Vector<float>(
network_states.GetLayer(input.handle).row(step_index - 1)),
input.previous_output, instance);
}
}
}
inline void MyelinDynamicComponentBase::BindOutputLayers(
size_t step_index, const NetworkStates &network_states,
sling::myelin::Instance *instance) const {
for (const OutputLayer &output : output_layers_) {
BindOutput(network_states.GetLayer(output.handle).row(step_index),
output.layer, instance);
}
}
inline FixedEmbeddings &MyelinDynamicComponentBase::GetFixedEmbeddings(
SessionState *session_state) const {
return session_state->extensions.Get(fixed_embeddings_handle_);
}
inline LinkedEmbeddings &MyelinDynamicComponentBase::GetLinkedEmbeddings(
SessionState *session_state) const {
return session_state->extensions.Get(linked_embeddings_handle_);
}
inline sling::myelin::Instance &MyelinDynamicComponentBase::GetInstance(
SessionState *session_state) const {
return session_state->extensions.Get(instance_handle_, cell_);
}
inline void MyelinDynamicComponentBase::MaybeTrace(
size_t step_index, sling::myelin::Instance *instance,
ComponentTrace *component_trace) const {
if (component_trace == nullptr) return;
while (component_trace->step_trace_size() <= step_index) {
component_trace->add_step_trace();
}
TraceMyelinInstance(instance,
component_trace->mutable_step_trace(step_index)
->AddExtension(CellTrace::step_trace_extension));
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELIN_DYNAMIC_COMPONENT_BASE_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include <memory>
#include <string>
#include "dragnn/core/test/generic.h"
#include "dragnn/protos/cell_trace.pb.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/protos/trace.pb.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/extensions.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/test/network_test_base.h"
#include "syntaxnet/base.h"
#include "sling/file/file.h"
#include "sling/myelin/flow.h"
#include <gmock/gmock.h>
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
using ::testing::_;
using ::testing::InSequence;
using ::testing::Invoke;
using ::testing::Return;
constexpr int kFlowVersion = 4;
constexpr int kVocabularySize = 123;
constexpr int kLogitsDim = 11;
constexpr int kNumSteps = 50;
class MyelinDynamicComponentTest : public NetworkTestBase {
protected:
// Options for building a Flow file for tests. By default, this specifies a
// working Flow file, but settings can be perturbed to trigger errors.
struct FlowFileOptions {
FlowFileOptions() = default;
// Name of the function to create.
string function_name = kTestComponentName;
// Dimension of the classification logits.
int logits_dim = kLogitsDim;
// Name of the variable containing the classification logits.
string logits_name = "logits";
// Type of the feature ID input.
sling::myelin::Type id_type = sling::myelin::DT_INT32;
// Dimension of the feature ID input.
int id_dim = 1;
};
// Builds and writes a simple Flow file. By default it produces a valid Flow,
// but arguments can be overridden for error testing. Returns the path to the
// Flow file.
static string WriteFlowFile() { return WriteFlowFile(FlowFileOptions()); }
static string WriteFlowFile(const FlowFileOptions &options) {
sling::myelin::Flow flow;
// A fixed feature ID input.
sling::myelin::Flow::Variable *id =
flow.AddVariable("id", options.id_type, {options.id_dim});
id->ref = true;
id->aliases.push_back(MakeMyelinInputFixedFeatureIdName(0, 0));
// An embedding matrix constant. Each embedding is filled with its index.
sling::myelin::Flow::Variable *embeddings =
flow.AddVariable("embeddings", sling::myelin::DT_FLOAT,
{kVocabularySize, options.logits_dim});
std::vector<float> data(kVocabularySize * options.logits_dim);
for (int row = 0; row < kVocabularySize; ++row) {
for (int column = 0; column < options.logits_dim; ++column) {
data[row * options.logits_dim + column] = row;
}
}
embeddings->SetData(data.data(), data.size() * sizeof(float));
// The retrieved embedding row, as logits.
sling::myelin::Flow::Variable *logits =
flow.AddVariable(options.logits_name, sling::myelin::DT_FLOAT,
{options.id_dim, options.logits_dim});
logits->ref = true;
logits->aliases.push_back(MakeMyelinOutputLayerName(options.logits_name));
// A Gather op that looks up the |id| in the |embeddings|, and returns the
// result in the |logits|.
flow.AddOperation(flow.AddFunction(options.function_name), "gather",
"Gather", {embeddings, id}, {logits});
const string flow_path =
tensorflow::io::JoinPath(tensorflow::testing::TmpDir(), "foo.flow");
sling::File::Init();
flow.Save(flow_path, kFlowVersion);
return flow_path;
}
// Creates a component, initializes it based on the |component_spec_text| and
// |flow_path|, and evaluates it. The |component_trace| is overwritten with
// traces, if non-null. On error, returns non-OK.
tensorflow::Status Run(const string &component_spec_text = "",
const string &flow_path = WriteFlowFile(),
ComponentTrace *component_trace = nullptr) {
ComponentSpec component_spec;
CHECK(TextFormat::ParseFromString(component_spec_text, &component_spec));
if (!component_spec.has_num_actions()) {
component_spec.set_num_actions(kLogitsDim);
}
component_spec.set_name(kTestComponentName);
auto *fixed_feature = component_spec.add_fixed_feature();
fixed_feature->set_embedding_dim(-1);
fixed_feature->set_size(1);
TF_RETURN_IF_ERROR(AddMyelinFlowResource(flow_path, &component_spec));
AddComponent(kTestComponentName);
TF_RETURN_IF_ERROR(
Component::CreateOrError("MyelinDynamicComponent", &component_));
TF_RETURN_IF_ERROR(component_->Initialize(component_spec, &variable_store_,
&network_state_manager_,
&extension_manager_));
network_states_.Reset(&network_state_manager_);
StartComponent(0); // MyelinDynamicComponent will add steps
session_state_.extensions.Reset(&extension_manager_);
TF_RETURN_IF_ERROR(component_->Evaluate(&session_state_, &compute_session_,
component_trace));
return tensorflow::Status::OK();
}
std::unique_ptr<Component> component_;
};
// Tests that MyelinDynamicComponent fails if the spec uses attention.
TEST_F(MyelinDynamicComponentTest, UnsupportedAttention) {
EXPECT_THAT(Run("attention_component:'foo'"),
test::IsErrorWithSubstr("Attention is not supported"));
}
// Tests that MyelinDynamicComponent fails if the spec has embedded fixed
// features.
TEST_F(MyelinDynamicComponentTest, InvalidFixedFeatureIsEmbedded) {
EXPECT_THAT(
Run("fixed_feature { embedding_dim:1 }"),
test::IsErrorWithSubstr("Myelin requires non-embedded fixed features"));
}
// Tests that MyelinDynamicComponent fails if the ComponentSpec has a fixed
// feature that does not appear in the Flow.
TEST_F(MyelinDynamicComponentTest, InvalidFixedFeatureNotInFlow) {
EXPECT_THAT(Run("fixed_feature { embedding_dim:-1 size:1 }"),
test::IsErrorWithSubstr(tensorflow::strings::StrCat(
"No Myelin tensor named '",
MakeMyelinInputFixedFeatureIdName(1, 0), "'")));
}
// Tests that MyelinDynamicComponent fails if the spec has multipled linked
// features.
TEST_F(MyelinDynamicComponentTest, InvalidLinkedFeatureIsMultiplied) {
EXPECT_THAT(Run("linked_feature { embedding_dim:1 }"),
test::IsErrorWithSubstr(
"Myelin requires non-multiplied linked features"));
}
// Tests that MyelinDynamicComponent fails if the ComponentSpec has a linked
// feature that does not appear in the Flow.
TEST_F(MyelinDynamicComponentTest, InvalidLinkedFeatureNotInFlow) {
const string kSpec = tensorflow::strings::StrCat(
"linked_feature { source_component:'", kTestComponentName,
"' source_layer:'logits' embedding_dim:-1 size:1 }");
EXPECT_THAT(Run(kSpec),
test::IsErrorWithSubstr(tensorflow::strings::StrCat(
"No Myelin tensor named '",
MakeMyelinInputLinkedActivationVectorName(0), "'")));
}
// Tests that MyelinDynamicComponent fails if the Flow file does not exist.
TEST_F(MyelinDynamicComponentTest, InvalidFlowFilePath) {
EXPECT_THAT(Run("", "/invalid/path"),
test::IsErrorWithSubstr("Failed to load Myelin Flow"));
}
// Tests that MyelinDynamicComponent fails if the function in the Flow file has
// the wrong name.
TEST_F(MyelinDynamicComponentTest, WrongFunctionName) {
FlowFileOptions options;
options.function_name = "wrong_function";
EXPECT_THAT(
Run("", WriteFlowFile(options)),
test::IsErrorWithSubstr(tensorflow::strings::StrCat(
"No function named '", kTestComponentName, "' in Myelin network")));
}
// Tests that MyelinDynamicComponent fails if the logits dimension does not
// match ComponentSpec.num_actions.
TEST_F(MyelinDynamicComponentTest, WrongLogitsDimension) {
FlowFileOptions options;
options.logits_dim = kLogitsDim + 1;
EXPECT_THAT(Run("", WriteFlowFile(options)),
test::IsErrorWithSubstr(
"Dimension mismatch between classification logits"));
}
// Tests that MyelinDynamicComponent fails if there is no "logits" layer.
TEST_F(MyelinDynamicComponentTest, WrongLogitsName) {
FlowFileOptions options;
options.logits_name = "not_logits";
EXPECT_THAT(Run("", WriteFlowFile(options)),
test::IsErrorWithSubstr("Unknown layer 'logits'"));
}
// Tests that MyelinDynamicComponent fails to compile if one of the Myelin
// tensors has the wrong type.
TEST_F(MyelinDynamicComponentTest, FailToCompile) {
FlowFileOptions options;
options.id_type = sling::myelin::DT_FLOAT;
EXPECT_THAT(Run("", WriteFlowFile(options)),
test::IsErrorWithSubstr("Failed to compile Myelin network"));
}
// Tests that MyelinDynamicComponent fails if one of the Myelin tensors is not
// vector-like.
TEST_F(MyelinDynamicComponentTest, NotVectorLike) {
FlowFileOptions options;
options.id_dim = 2;
EXPECT_THAT(
Run("", WriteFlowFile(options)),
test::IsErrorWithSubstr("Myelin tensor has non-vector-like shape"));
}
// Tests that MyelinDynamicComponent fails if AdvanceFromPrediction() fails.
TEST_F(MyelinDynamicComponentTest, FailToAdvanceFromPrediction) {
EXPECT_CALL(compute_session_, IsTerminal(_)).WillRepeatedly(Return(false));
EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
.WillOnce(Return(false));
EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
.WillOnce(Invoke(ExtractFeatures(0, {{10, 1.0}})));
EXPECT_THAT(Run(), test::IsErrorWithSubstr(
"Error in ComputeSession::AdvanceFromPrediction()"));
}
// Tests that MyelinDynamicComponent can run a simple non-deterministic Flow.
TEST_F(MyelinDynamicComponentTest, SimpleNonDeterministicFlow) {
SetupTransitionLoop(kNumSteps);
EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
.Times(kNumSteps)
.WillRepeatedly(Return(true));
{ // Extract a sequence of feature IDs equal to 2 * step_index.
ASSERT_LE(2 * kNumSteps, kVocabularySize);
InSequence scoped;
for (int step_index = 0; step_index < kNumSteps; ++step_index) {
EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
.WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
}
}
TF_ASSERT_OK(Run());
const Matrix<float> logits(GetLayer(kTestComponentName, "logits"));
ASSERT_EQ(logits.num_rows(), kNumSteps);
ASSERT_EQ(logits.num_columns(), kLogitsDim);
// Since each row of the embedding matrix is filled with its index, the logits
// should be equal to the feature IDs.
for (int step_index = 0; step_index < kNumSteps; ++step_index) {
ExpectVector(logits.row(step_index), kLogitsDim, 2 * step_index);
}
}
// Tests that MyelinDynamicComponent can run a simple deterministic Flow.
TEST_F(MyelinDynamicComponentTest, SimpleDeterministicFlow) {
SetupTransitionLoop(kNumSteps);
EXPECT_CALL(compute_session_, AdvanceFromOracle(kTestComponentName))
.Times(kNumSteps);
{ // Extract a sequence of feature IDs equal to 2 * step_index.
ASSERT_LE(2 * kNumSteps, kVocabularySize);
InSequence scoped;
for (int step_index = 0; step_index < kNumSteps; ++step_index) {
EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
.WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
}
}
FlowFileOptions options;
options.logits_dim = 1;
TF_ASSERT_OK(Run("num_actions:1", WriteFlowFile(options)));
}
// Tests that MyelinDynamicComponent can run a simple Flow with tracing enabled.
TEST_F(MyelinDynamicComponentTest, SimpleFlowWithTracing) {
SetupTransitionLoop(kNumSteps);
EXPECT_CALL(compute_session_, AdvanceFromPrediction(_, _, _, _))
.Times(kNumSteps)
.WillRepeatedly(Return(true));
{ // Extract a sequence of feature IDs equal to 2 * step_index.
ASSERT_LE(2 * kNumSteps, kVocabularySize);
InSequence scoped;
for (int step_index = 0; step_index < kNumSteps; ++step_index) {
EXPECT_CALL(compute_session_, GetInputFeatures(_, _, _, _, _))
.WillOnce(Invoke(ExtractFeatures(0, {{2 * step_index, 1.0}})));
}
}
ComponentTrace component_trace;
TF_ASSERT_OK(Run("", WriteFlowFile(), &component_trace));
// Each step trace should have a cell trace from the Myelin instance.
ASSERT_EQ(component_trace.step_trace_size(), kNumSteps);
for (const ComponentStepTrace &step_trace : component_trace.step_trace()) {
ASSERT_EQ(step_trace.ExtensionSize(CellTrace::step_trace_extension), 1);
const CellTrace &cell_trace =
step_trace.GetExtension(CellTrace::step_trace_extension, 0);
EXPECT_EQ(cell_trace.name(), kTestComponentName);
}
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_library.h"
#include <string>
#include "syntaxnet/base.h"
#include "tensorflow/core/lib/strings/strcat.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
bool PreMultipliedEmbeddings::Transform(sling::myelin::Flow *flow) {
bool transformed_something = false;
for (sling::myelin::Flow::Operation *matmul :
flow->Find({"Gather", "MatMul"})) {
if (matmul->indegree() != 2) continue;
sling::myelin::Flow::Variable *gathered = matmul->inputs[0];
sling::myelin::Flow::Variable *weights = matmul->inputs[1];
sling::myelin::Flow::Operation *gather = gathered->producer;
if (gather->indegree() != 2) continue;
sling::myelin::Flow::Variable *embeddings = gather->inputs[0];
sling::myelin::Flow::Variable *indices = gather->inputs[1];
if (gathered->out) continue;
if (!weights->constant()) continue;
if (weights->rank() != 2) continue;
if (!embeddings->constant()) continue;
if (embeddings->rank() != 2) continue;
if (embeddings->type != weights->type) continue;
// Add an operation to pre-multiply the embeddings and weights.
const string product_name =
tensorflow::strings::StrCat(embeddings->name, "/", weights->name);
const string pre_multiply_name =
tensorflow::strings::StrCat(product_name, "/PreMultiply");
sling::myelin::Flow::Variable *product = flow->AddVariable(
product_name, weights->type, {embeddings->dim(0), weights->dim(1)});
flow->AddOperation(gather->func, pre_multiply_name, "MatMul",
{embeddings, weights}, {product});
// Convert the MatMul into a Gather on the pre-multiplied embeddings.
matmul->type = "Gather";
matmul->ReplaceInput(gathered, product);
matmul->ReplaceInput(weights, indices);
// Remove the original Gather if it is no longer used.
if (gathered->consumers.empty()) flow->RemoveOperation(gather);
transformed_something = true;
}
return transformed_something;
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
// Myelin typers, transformers, and kernels specific to the DRAGNN runtime.
#ifndef DRAGNN_RUNTIME_MYELIN_MYELIN_LIBRARY_H_
#define DRAGNN_RUNTIME_MYELIN_MYELIN_LIBRARY_H_
#include "sling/myelin/flow.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// Rearranges the flow to allow the "pre-multiplied embeddings" optimization.
// Specifically, performs the following transformation:
//
// tf.matmul(tf.gather(embeddings, indices), weights) =
// tf.gather(tf.matmul(embeddings, weights), indices)
//
// The transformation only applies if the embeddings and weights are constants.
// Myelin has constant folding transformations that will trigger and pre-compute
// the multiplication of the embeddings and weights.
//
// NB: There is already a PrecomputedEmbeddings transformer in Myelin but that
// operates on the Lookup op and expects an intervening Reshape.
class PreMultipliedEmbeddings : public sling::myelin::Transformer {
public:
// Implements Transformer.
bool Transform(sling::myelin::Flow *flow) override;
};
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELIN_LIBRARY_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_library.h"
#include <vector>
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Tests that PreMultipliedEmbeddings does nothing on an empty Flow.
TEST(PreMultipliedEmbeddingsTest, DoesNothingOnEmptyFlow) {
sling::myelin::Flow flow;
PreMultipliedEmbeddings transformer;
EXPECT_FALSE(transformer.Transform(&flow));
}
// Tests that PreMultipliedEmbeddings can rearrange a MatMul of a Gather into a
// Gather of a pre-multiplied matrix.
TEST(PreMultipliedEmbeddingsTest, AppliesPreMultiplication) {
sling::myelin::Flow flow;
sling::myelin::Flow::Function *function = flow.AddFunction("test_function");
sling::myelin::Flow::Variable *indices =
flow.AddVariable("indices", sling::myelin::DT_INT32, {1});
sling::myelin::Flow::Variable *embeddings =
flow.AddVariable("embeddings", sling::myelin::DT_FLOAT, {10, 20});
sling::myelin::Flow::Variable *gathered =
flow.AddVariable("gathered", sling::myelin::DT_FLOAT, {1, 20});
sling::myelin::Flow::Variable *weights =
flow.AddVariable("weights", sling::myelin::DT_FLOAT, {20, 30});
sling::myelin::Flow::Variable *output =
flow.AddVariable("output", sling::myelin::DT_FLOAT, {1, 30});
flow.AddOperation(function, "gather", "Gather", {embeddings, indices},
{gathered});
flow.AddOperation(function, "matmul", "MatMul", {gathered, weights},
{output});
// Attach constant data to the matrices.
const std::vector<float> floats(20 * 30); // big enough for both
embeddings->SetData(floats.data(), 10 * 20 * sizeof(float));
weights->SetData(floats.data(), 20 * 30 * sizeof(float));
PreMultipliedEmbeddings transformer;
ASSERT_TRUE(transformer.Transform(&flow));
sling::myelin::Flow::Variable *product = flow.Var("embeddings/weights");
ASSERT_NE(product, nullptr);
ASSERT_EQ(product->rank(), 2);
EXPECT_EQ(product->dim(0), 10);
EXPECT_EQ(product->dim(1), 30);
sling::myelin::Flow::Operation *pre_multiply =
flow.Op("embeddings/weights/PreMultiply");
ASSERT_NE(pre_multiply, nullptr);
ASSERT_EQ(pre_multiply->indegree(), 2);
ASSERT_EQ(pre_multiply->outdegree(), 1);
EXPECT_EQ(pre_multiply->type, "MatMul");
EXPECT_EQ(pre_multiply->inputs[0], embeddings);
EXPECT_EQ(pre_multiply->inputs[1], weights);
EXPECT_EQ(pre_multiply->outputs[0], product);
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include <algorithm>
#include "dragnn/runtime/myelin/myelin_library.h"
#include "sling/base/status.h"
#include "sling/file/file.h"
#include "sling/myelin/kernel/tensorflow.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
const char *const kMyelinFlowResourceName = "myelin-flow";
const char *const kMyelinFlowResourceFileFormat = "model";
const char *const kMyelinFlowResourceRecordFormat = "sling.myelin.Flow";
tensorflow::Status LookupMyelinFlowResource(const ComponentSpec &component_spec,
const Resource **flow_resource) {
const Resource *found_resource = nullptr;
for (const Resource &resource : component_spec.resource()) {
if (resource.name() != kMyelinFlowResourceName) continue;
if (found_resource != nullptr) {
return tensorflow::errors::InvalidArgument(
"Component '", component_spec.name(),
"' contains duplicate Myelin Flow resources");
}
if (resource.part_size() != 1) {
return tensorflow::errors::InvalidArgument(
"Component '", component_spec.name(),
"' has malformed Myelin Flow resource; expected 1 part");
}
const Part &part = resource.part(0);
if (part.file_format() != kMyelinFlowResourceFileFormat) {
return tensorflow::errors::InvalidArgument(
"Component '", component_spec.name(),
"' has malformed Myelin Flow resource; wrong file format");
}
if (part.record_format() != kMyelinFlowResourceRecordFormat) {
return tensorflow::errors::InvalidArgument(
"Component '", component_spec.name(),
"' has malformed Myelin Flow resource; wrong record format");
}
found_resource = &resource;
}
if (found_resource == nullptr) {
return tensorflow::errors::NotFound("Component '", component_spec.name(),
"' has no Myelin Flow resource");
}
// Success; make modifications.
*flow_resource = found_resource;
return tensorflow::Status::OK();
}
tensorflow::Status AddMyelinFlowResource(const string &path,
ComponentSpec *component_spec) {
if (std::any_of(component_spec->resource().begin(),
component_spec->resource().end(),
[](const Resource &resource) {
return resource.name() == kMyelinFlowResourceName;
})) {
return tensorflow::errors::InvalidArgument(
"Component '", component_spec->name(),
"' already contains a Myelin Flow resource");
}
// Success; make modifications.
Resource *resource = component_spec->add_resource();
resource->set_name(kMyelinFlowResourceName);
Part *part = resource->add_part();
part->set_file_pattern(path);
part->set_file_format(kMyelinFlowResourceFileFormat);
part->set_record_format(kMyelinFlowResourceRecordFormat);
return tensorflow::Status::OK();
}
tensorflow::Status LoadMyelinFlow(const string &flow_path,
sling::myelin::Flow *flow) {
sling::File::Init();
const sling::Status status = flow->Load(flow_path);
if (!status.ok()) {
return tensorflow::errors::Internal("Failed to load Myelin Flow from '",
flow_path, ": ", status.ToString());
}
// Mark cell inputs and outputs.
for (sling::myelin::Flow::Variable *variable : flow->vars()) {
for (tensorflow::StringPiece alias : variable->aliases) {
if (tensorflow::str_util::StartsWith(alias, "INPUT/")) {
variable->in = true;
}
if (tensorflow::str_util::StartsWith(alias, "OUTPUT/")) {
variable->out = true;
}
}
}
return tensorflow::Status::OK();
}
void RegisterMyelinLibraries(sling::myelin::Library *library) {
// TODO(googleuser): Add more libraries?
sling::myelin::RegisterTensorflowLibrary(library);
library->RegisterTransformer(new PreMultipliedEmbeddings());
}
std::set<string> GetRecurrentLayerNames(const sling::myelin::Flow &flow) {
std::set<string> names;
for (const sling::myelin::Flow::Variable *variable : flow.vars()) {
for (tensorflow::StringPiece alias : variable->aliases) {
if (!tensorflow::str_util::ConsumePrefix(&alias, "INPUT/")) continue;
if (tensorflow::str_util::ConsumePrefix(&alias, "fixed_channel_")) {
continue;
}
if (tensorflow::str_util::ConsumePrefix(&alias, "linked_channel_")) {
continue;
}
names.insert(alias.ToString());
}
}
return names;
}
std::set<string> GetOutputLayerNames(const sling::myelin::Flow &flow) {
std::set<string> names;
for (const sling::myelin::Flow::Variable *variable : flow.vars()) {
for (tensorflow::StringPiece alias : variable->aliases) {
if (!tensorflow::str_util::ConsumePrefix(&alias, "OUTPUT/")) continue;
names.insert(alias.ToString());
}
}
return names;
}
string MakeMyelinInputFixedFeatureIdName(int channel_id, int index) {
return tensorflow::strings::StrCat(
"INPUT/fixed_channel_", channel_id, "_index_", index, "_ids");
}
string MakeMyelinInputLinkedActivationVectorName(int channel_id) {
return tensorflow::strings::StrCat("INPUT/linked_channel_", channel_id,
"_activations");
}
string MakeMyelinInputLinkedOutOfBoundsIndicatorName(int channel_id) {
return tensorflow::strings::StrCat("INPUT/linked_channel_", channel_id,
"_out_of_bounds");
}
string MakeMyelinInputRecurrentLayerName(const string &layer_name) {
return tensorflow::strings::StrCat("INPUT/", layer_name);
}
string MakeMyelinOutputLayerName(const string &layer_name) {
return tensorflow::strings::StrCat("OUTPUT/", layer_name);
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
// Utils for working with specifications of Myelin-based DRAGNN runtime models.
#ifndef DRAGNN_RUNTIME_MYELIN_MYELIN_SPEC_UTILS_H_
#define DRAGNN_RUNTIME_MYELIN_MYELIN_SPEC_UTILS_H_
#include <set>
#include <string>
#include "dragnn/protos/spec.pb.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "sling/myelin/flow.h"
#include "tensorflow/core/lib/core/status.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// The name, file format, and record format of the resource that contains the
// Myelin Flow for each component.
extern const char *const kMyelinFlowResourceName;
extern const char *const kMyelinFlowResourceFileFormat;
extern const char *const kMyelinFlowResourceRecordFormat;
// Points |flow_resource| to the resource in the |component_spec| that specifies
// the Myelin Flow file. On error, returns non-OK and modifies nothing.
tensorflow::Status LookupMyelinFlowResource(const ComponentSpec &component_spec,
const Resource **flow_resource);
// Adds a resource to the |component_spec| that specifies the Myelin Flow file
// at the |path|. On error, returns non-OK and modifies nothing.
tensorflow::Status AddMyelinFlowResource(const string &path,
ComponentSpec *component_spec);
// Loads a Myelin Flow file from the |flow_path| into the |flow| and ensures
// that inputs and outputs are marked properly. On error, returns non-OK.
tensorflow::Status LoadMyelinFlow(const string &flow_path,
sling::myelin::Flow *flow);
// Registers a standard set of libraries in the Myelin |library|.
void RegisterMyelinLibraries(sling::myelin::Library *library);
// Returns the set of recurrent input layer names in the |flow|. A recurrent
// input layer is defined as any input that is not a fixed or linked feature.
//
// Note that recurrent input layers differ from recurrent linked features. The
// latter are linked features that have been configured to refer to the current
// component, while the former are hard-coded in the network structure itself.
// See, for example, the context tensor arrays that hold the cell state in the
// LstmNetwork.
//
// TODO(googleuser): Use a more robust naming scheme for recurrent inputs?
std::set<string> GetRecurrentLayerNames(const sling::myelin::Flow &flow);
// Returns the set of output layer names in the |flow|.
std::set<string> GetOutputLayerNames(const sling::myelin::Flow &flow);
// Returns the name of the Myelin input for the ID of the |index|'th feature in
// the |channel_id|'th fixed feature channel.
string MakeMyelinInputFixedFeatureIdName(int channel_id, int index);
// Returns the names of the Myelin inputs for the source activation vector and
// out-of-bounds indicator of the |channel_id|'th linked feature channel.
string MakeMyelinInputLinkedActivationVectorName(int channel_id);
string MakeMyelinInputLinkedOutOfBoundsIndicatorName(int channel_id);
// Returns the name of the Myelin input for the hard-coded recurrent layer named
// |layer_name|.
string MakeMyelinInputRecurrentLayerName(const string &layer_name);
// Returns the name of the Myelin output for the layer named |layer_name|.
string MakeMyelinOutputLayerName(const string &layer_name);
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELIN_SPEC_UTILS_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include <set>
#include <string>
#include "dragnn/core/test/generic.h"
#include "dragnn/protos/spec.pb.h"
#include "syntaxnet/base.h"
#include "sling/file/file.h"
#include "sling/myelin/compute.h"
#include "sling/myelin/flow.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
TEST(MyelinSpecUtilsTest, AddAndLookupMyelinFlowResource) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
const Resource *resource = nullptr;
TF_ASSERT_OK(LookupMyelinFlowResource(component_spec, &resource));
ASSERT_NE(resource, nullptr);
EXPECT_EQ(resource->name(), kMyelinFlowResourceName);
ASSERT_EQ(resource->part_size(), 1);
EXPECT_EQ(resource->part(0).file_pattern(), "/dev/null");
EXPECT_EQ(resource->part(0).file_format(), kMyelinFlowResourceFileFormat);
EXPECT_EQ(resource->part(0).record_format(), kMyelinFlowResourceRecordFormat);
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceMissing) {
ComponentSpec component_spec;
const Resource *resource = nullptr;
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("has no Myelin Flow resource"));
component_spec.add_resource()->set_name("foo");
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("has no Myelin Flow resource"));
component_spec.add_resource()->set_name("bar");
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("has no Myelin Flow resource"));
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceWrongName) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
component_spec.mutable_resource(0)->set_name("bad");
const Resource *resource = nullptr;
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("has no Myelin Flow resource"));
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceWrongFileFormat) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
component_spec.mutable_resource(0)->mutable_part(0)->set_file_format("bad");
const Resource *resource = nullptr;
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("wrong file format"));
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceWrongRecordFormat) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
component_spec.mutable_resource(0)->mutable_part(0)->set_record_format("bad");
const Resource *resource = nullptr;
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("wrong record format"));
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceWrongNumberOfParts) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
component_spec.mutable_resource(0)->add_part();
const Resource *resource = nullptr;
EXPECT_THAT(LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("expected 1 part"));
}
TEST(MyelinSpecUtilsTest, LookupMyelinFlowResourceDuplicate) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
component_spec.add_resource()->set_name(kMyelinFlowResourceName);
const Resource *resource = nullptr;
EXPECT_THAT(
LookupMyelinFlowResource(component_spec, &resource),
test::IsErrorWithSubstr("contains duplicate Myelin Flow resource"));
}
TEST(MyelinSpecUtilsTest, AddMyelinFlowResourceDuplicate) {
ComponentSpec component_spec;
TF_ASSERT_OK(AddMyelinFlowResource("/dev/null", &component_spec));
EXPECT_THAT(
AddMyelinFlowResource("another/flow", &component_spec),
test::IsErrorWithSubstr("already contains a Myelin Flow resource"));
}
TEST(MyelinSpecUtilsTest, LoadMyelinFlowInvalidPath) {
sling::myelin::Flow flow;
EXPECT_THAT(LoadMyelinFlow("invalid/path", &flow),
test::IsErrorWithSubstr("Failed to load Myelin Flow"));
}
TEST(MyelinSpecUtilsTest, LoadMyelinFlowValidFile) {
// Build and write a Flow file with some variables that are annotated with
// input and output aliases.
sling::myelin::Flow original_flow;
original_flow
.AddVariable("input", sling::myelin::DT_FLOAT, sling::myelin::Shape())
->aliases = {"INPUT/a"};
original_flow
.AddVariable("output", sling::myelin::DT_FLOAT, sling::myelin::Shape())
->aliases = {"OUTPUT/b"};
original_flow
.AddVariable("both", sling::myelin::DT_FLOAT, sling::myelin::Shape())
->aliases = {"INPUT/c", "OUTPUT/d"};
original_flow.AddVariable("neither", sling::myelin::DT_FLOAT,
sling::myelin::Shape());
const string flow_path =
tensorflow::io::JoinPath(tensorflow::testing::TmpDir(), "foo.flow");
sling::File::Init();
original_flow.Save(flow_path);
// Load the Flow file into a fresh Flow and check that inputs and outputs are
// marked as such.
sling::myelin::Flow flow;
TF_ASSERT_OK(LoadMyelinFlow(flow_path, &flow));
ASSERT_NE(flow.Var("input"), nullptr);
EXPECT_TRUE(flow.Var("input")->in);
EXPECT_FALSE(flow.Var("input")->out);
ASSERT_NE(flow.Var("output"), nullptr);
EXPECT_FALSE(flow.Var("output")->in);
EXPECT_TRUE(flow.Var("output")->out);
ASSERT_NE(flow.Var("both"), nullptr);
EXPECT_TRUE(flow.Var("both")->in);
EXPECT_TRUE(flow.Var("both")->out);
ASSERT_NE(flow.Var("neither"), nullptr);
EXPECT_FALSE(flow.Var("neither")->in);
EXPECT_FALSE(flow.Var("neither")->out);
}
TEST(MyelinSpecUtilsTest, RegisterMyelinLibraries) {
sling::myelin::Library library;
RegisterMyelinLibraries(&library);
// The |library| should contain something.
EXPECT_GT(library.transformers().size() + library.typers().size(), 0);
}
TEST(MyelinSpecUtilsTest, GetRecurrentLayerNamesEmpty) {
sling::myelin::Flow flow;
const std::set<string> expected_names;
EXPECT_EQ(GetRecurrentLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetRecurrentLayerNamesVariablesWithNoAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {});
flow.AddVariable("y", sling::myelin::DT_INT32, {});
const std::set<string> expected_names;
EXPECT_EQ(GetRecurrentLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetRecurrentLayerNamesVariablesWithAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {})->aliases = {"foo", "bar"};
flow.AddVariable("y", sling::myelin::DT_INT32, {})->aliases = {
"INPUT/y", //
"INPUT/fixed_channel_0_index_0_ids", //
"INPUT/linked_channel_0_activations"};
flow.AddVariable("z", sling::myelin::DT_INT32, {})->aliases = {"OUTPUT/z"};
const std::set<string> expected_names = {"y"};
EXPECT_EQ(GetRecurrentLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetRecurrentLayerNamesVariablesWithMultipleAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {})->aliases = {"foo", "bar"};
flow.AddVariable("y", sling::myelin::DT_INT32, {})->aliases = {
"INPUT/recurrent_1", //
"INPUT/recurrent_2", //
"INPUT/fixed_channel_0_index_0_ids", //
"INPUT/linked_channel_0_activations"};
flow.AddVariable("z", sling::myelin::DT_INT32, {})->aliases = {
"OUTPUT/output_1", //
"OUTPUT/output_2"};
const std::set<string> expected_names = {"recurrent_1", "recurrent_2"};
EXPECT_EQ(GetRecurrentLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetOutputLayerNamesEmpty) {
sling::myelin::Flow flow;
const std::set<string> expected_names;
EXPECT_EQ(GetOutputLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetOutputLayerNamesVariablesWithNoAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {});
flow.AddVariable("y", sling::myelin::DT_INT32, {});
const std::set<string> expected_names;
EXPECT_EQ(GetOutputLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetOutputLayerNamesVariablesWithAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {})->aliases = {"foo", "bar"};
flow.AddVariable("y", sling::myelin::DT_INT32, {})->aliases = {
"INPUT/y", //
"INPUT/fixed_channel_0_index_0_ids", //
"INPUT/linked_channel_0_activations"};
flow.AddVariable("z", sling::myelin::DT_INT32, {})->aliases = {"OUTPUT/z"};
const std::set<string> expected_names = {"z"};
EXPECT_EQ(GetOutputLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, GetOutputLayerNamesVariablesWithMultipleAliases) {
sling::myelin::Flow flow;
flow.AddVariable("x", sling::myelin::DT_FLOAT, {})->aliases = {"foo", "bar"};
flow.AddVariable("y", sling::myelin::DT_INT32, {})->aliases = {
"INPUT/recurrent_1", //
"INPUT/recurrent_2", //
"INPUT/fixed_channel_0_index_0_ids", //
"INPUT/linked_channel_0_activations"};
flow.AddVariable("z", sling::myelin::DT_INT32, {})->aliases = {
"OUTPUT/output_1", //
"OUTPUT/output_2"};
const std::set<string> expected_names = {"output_1", "output_2"};
EXPECT_EQ(GetOutputLayerNames(flow), expected_names);
}
TEST(MyelinSpecUtilsTest, MakeMyelinInputFixedFeatureIdName) {
EXPECT_EQ(MakeMyelinInputFixedFeatureIdName(0, 1),
"INPUT/fixed_channel_0_index_1_ids");
EXPECT_EQ(MakeMyelinInputFixedFeatureIdName(1, 0),
"INPUT/fixed_channel_1_index_0_ids");
}
TEST(MyelinSpecUtilsTest, MakeMyelinInputLinkedActivationVectorName) {
EXPECT_EQ(MakeMyelinInputLinkedActivationVectorName(0),
"INPUT/linked_channel_0_activations");
EXPECT_EQ(MakeMyelinInputLinkedActivationVectorName(1),
"INPUT/linked_channel_1_activations");
}
TEST(MyelinSpecUtilsTest, MakeMyelinInputLinkedOutOfBoundsIndicatorName) {
EXPECT_EQ(MakeMyelinInputLinkedOutOfBoundsIndicatorName(0),
"INPUT/linked_channel_0_out_of_bounds");
EXPECT_EQ(MakeMyelinInputLinkedOutOfBoundsIndicatorName(1),
"INPUT/linked_channel_1_out_of_bounds");
}
TEST(MyelinSpecUtilsTest, MakeMyelinInputRecurrentLayerName) {
EXPECT_EQ(MakeMyelinInputRecurrentLayerName("foo"), "INPUT/foo");
EXPECT_EQ(MakeMyelinInputRecurrentLayerName("bar_baz"), "INPUT/bar_baz");
}
TEST(MyelinSpecUtilsTest, MakeMyelinOutputLayerName) {
EXPECT_EQ(MakeMyelinOutputLayerName("foo"), "OUTPUT/foo");
EXPECT_EQ(MakeMyelinOutputLayerName("bar_baz"), "OUTPUT/bar_baz");
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_tracing.h"
#include <map>
#include <string>
#include "syntaxnet/base.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Copies |num_values| |T|s from |data| into the |tensor_trace|. If |T| does
// not match the |type|, returns false and modifies nothing. The bool return
// allows this function to be chained until a matching type is found.
template <class T>
bool TryCopyValues(sling::myelin::Type type, const char *data, int num_values,
CellTensorTrace *tensor_trace) {
if (sling::myelin::Traits<T>().type() != type) return false;
const T *begin = reinterpret_cast<const T *>(data);
const T *end = begin + num_values;
tensor_trace->clear_value();
for (; begin != end; ++begin) tensor_trace->add_value(*begin);
return true;
}
} // namespace
void TraceMyelinInstance(sling::myelin::Instance *instance,
CellTrace *cell_trace) {
const sling::myelin::Cell &cell = *instance->cell();
cell_trace->Clear();
cell_trace->set_name(cell.name());
// Collect steps and tensors in sorted maps for deterministic ordering.
std::map<string, const sling::myelin::Step *> steps;
std::map<string, sling::myelin::Tensor *> tensors;
for (const sling::myelin::Step *step : cell.steps()) {
steps[step->name()] = step;
for (sling::myelin::Tensor *tensor : step->inputs()) {
tensors[tensor->name()] = tensor;
}
for (sling::myelin::Tensor *tensor : step->outputs()) {
tensors[tensor->name()] = tensor;
}
}
// Trace each step as an operation.
for (const auto &it : steps) {
const sling::myelin::Step *step = it.second;
CellOperationTrace *operation_trace = cell_trace->add_operation();
operation_trace->set_name(step->name());
operation_trace->set_type(step->type());
operation_trace->set_kernel(step->kernel()->Name());
for (sling::myelin::Tensor *tensor : step->inputs()) {
operation_trace->add_input(tensor->name());
}
for (sling::myelin::Tensor *tensor : step->outputs()) {
operation_trace->add_output(tensor->name());
}
}
// Trace each tensor and its value.
for (const auto &it : tensors) {
sling::myelin::Tensor *tensor = it.second;
if (!tensor->IsLocal()) continue; // ignore globals; e.g., weight matrices
const string &name = tensor->name();
const sling::myelin::Type type = tensor->type();
// Find the variable data for the |tensor|. Note that ref tensors need to
// be dereferenced.
const char *data = instance->GetAddress(tensor);
if (tensor->ref()) data = *reinterpret_cast<const char *const *>(data);
const int size = tensor->aligned().elements();
CellTensorTrace *tensor_trace = cell_trace->add_tensor();
tensor_trace->set_name(name);
tensor_trace->set_type(sling::myelin::TypeTraits::of(type).name());
for (int i = 0; i < tensor->rank(); ++i) {
tensor_trace->add_dimension(tensor->dim(i));
tensor_trace->add_aligned_dimension(tensor->aligned(i));
}
switch (tensor->order()) {
case sling::myelin::ROW_MAJOR:
tensor_trace->set_order(CellTensorTrace::ORDER_ROW_MAJOR);
break;
case sling::myelin::COLUMN_MAJOR:
tensor_trace->set_order(CellTensorTrace::ORDER_COLUMN_MAJOR);
break;
default:
break;
}
// Try copying tensor data using all relevant types. At most one attempt
// will succeed and modify the |tensor_trace|.
if (!TryCopyValues<float>(type, data, size, tensor_trace) &&
!TryCopyValues<double>(type, data, size, tensor_trace) &&
!TryCopyValues<bool>(type, data, size, tensor_trace) &&
!TryCopyValues<int8>(type, data, size, tensor_trace) &&
!TryCopyValues<int16>(type, data, size, tensor_trace) &&
!TryCopyValues<int32>(type, data, size, tensor_trace) &&
!TryCopyValues<int64>(type, data, size, tensor_trace) &&
!TryCopyValues<uint8>(type, data, size, tensor_trace) &&
!TryCopyValues<uint16>(type, data, size, tensor_trace)) {
LOG(WARNING) << "Can't convert data for tensor " << name << " with type "
<< tensor_trace->type();
}
}
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#ifndef DRAGNN_RUNTIME_MYELIN_MYELIN_TRACING_H_
#define DRAGNN_RUNTIME_MYELIN_MYELIN_TRACING_H_
#include "dragnn/protos/cell_trace.pb.h"
#include "sling/myelin/compute.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// Overwrites the |cell_trace| with traces extracted from the |instance|. Does
// not modify the |instance|; it is non-const because the relevant accessors are
// declared non-const.
void TraceMyelinInstance(sling::myelin::Instance *instance,
CellTrace *cell_trace);
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELIN_TRACING_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelin_tracing.h"
#include <string.h>
#include <string>
#include "dragnn/core/test/generic.h"
#include "dragnn/protos/cell_trace.pb.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/test/helpers.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "sling/myelin/flow.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/dynamic_annotations.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Name of the dummy cell for tests.
constexpr char kCellName[] = "test_cell";
// Returns a CellTrace parsed from the concatenation of the |args|.
template <class... Args>
CellTrace ParseCellTrace(const Args &... args) {
const string text_proto = tensorflow::strings::StrCat(args...);
CellTrace cell_trace;
CHECK(TextFormat::ParseFromString(text_proto, &cell_trace));
return cell_trace;
}
// Testing rig.
class TraceMyelinInstanceTest : public ::testing::Test {
protected:
// Compiles the |flow_|, binds the name=>data |feeds|, evaluates the cell, and
// returns an extracted trace.
CellTrace GetTrace(const std::map<string, MutableAlignedView> &feeds) {
sling::myelin::Library library;
RegisterMyelinLibraries(&library);
LOG(INFO) << "Original flow:\n" << flow_.ToString();
flow_.Analyze(library);
LOG(INFO) << "Analyzed flow:\n" << flow_.ToString();
sling::myelin::Network network;
CHECK(network.Compile(flow_, library));
const sling::myelin::Cell *cell = network.GetCell(kCellName);
CHECK(cell != nullptr) << "Unknown cell: " << kCellName;
sling::myelin::Instance instance(cell);
for (const auto &it : feeds) {
const string &name = it.first;
char *data = it.second.data();
sling::myelin::Tensor *tensor = network.GetParameter(name);
CHECK(tensor != nullptr) << "Unknown tensor: " << name;
instance.SetReference(tensor, data);
}
instance.Compute();
CellTrace cell_trace;
TraceMyelinInstance(&instance, &cell_trace);
return cell_trace;
}
// Flow, to be modified in each test.
sling::myelin::Flow flow_;
// The function to trace. Each test should add operations to this.
sling::myelin::Flow::Function *function_ = flow_.AddFunction(kCellName);
};
// Tests tracing on a simple cell with one operation. In this cell, both the
// input and output are Tensor refs and need to be fed.
TEST_F(TraceMyelinInstanceTest, SingleOperation) {
sling::myelin::Flow::Variable *input =
flow_.AddVariable("input", sling::myelin::DT_FLOAT, {1});
input->in = true;
input->ref = true;
sling::myelin::Flow::Variable *one =
flow_.AddVariable("one", sling::myelin::DT_FLOAT, {1});
constexpr float kOne = 1.0;
one->SetData(&kOne, sizeof(float));
sling::myelin::Flow::Variable *axis =
flow_.AddVariable("axis", sling::myelin::DT_INT32, {1});
constexpr int32 kAxis = 0;
axis->SetData(&kAxis, sizeof(int32));
sling::myelin::Flow::Variable *output =
flow_.AddVariable("output", sling::myelin::DT_FLOAT, {2});
output->out = true;
output->ref = true;
sling::myelin::Flow::Operation *concat = flow_.AddOperation(
function_, "concat", "ConcatV2", {input, one, axis}, {output});
concat->SetAttr("N", 2);
UniqueVector<float> input_feed(1);
UniqueVector<float> output_feed(2);
(*input_feed)[0] = -1.5;
TF_ANNOTATE_MEMORY_IS_INITIALIZED(output_feed->data(),
output_feed->size() * sizeof(float));
const std::map<string, MutableAlignedView> feeds = {
{"input", input_feed.view()}, //
{"output", output_feed.view()}};
const CellTrace expected_trace = ParseCellTrace(R"(
name: ')", kCellName, R"('
tensor {
name: 'input'
type: 'float32'
dimension: [1]
aligned_dimension: [1]
order: ORDER_ROW_MAJOR
value: [-1.5]
}
tensor {
name: 'output'
type: 'float32'
dimension: [2]
aligned_dimension: [2]
order: ORDER_ROW_MAJOR
value: [-1.5, 1.0]
}
operation {
name: 'concat'
type: 'ConcatV2'
kernel: 'BasicConcat'
input: ['input', 'one', 'axis']
output: ['output']
}
)");
EXPECT_THAT(GetTrace(feeds), test::EqualsProto(expected_trace));
EXPECT_EQ((*output_feed)[0], -1.5);
EXPECT_EQ((*output_feed)[1], 1.0);
}
// Tests tracing on a slightly more complex cell with a few operations. In this
// case, only the input is a Tensor ref and needs to be fed.
TEST_F(TraceMyelinInstanceTest, MultiOperation) {
sling::myelin::Flow::Variable *input =
flow_.AddVariable("input", sling::myelin::DT_FLOAT, {1});
input->in = true;
input->ref = true;
sling::myelin::Flow::Variable *one =
flow_.AddVariable("one", sling::myelin::DT_FLOAT, {1});
constexpr float kOne = 1.0;
one->SetData(&kOne, sizeof(float));
sling::myelin::Flow::Variable *two =
flow_.AddVariable("two", sling::myelin::DT_FLOAT, {1});
constexpr float kTwo = 2.0;
two->SetData(&kTwo, sizeof(float));
sling::myelin::Flow::Variable *three =
flow_.AddVariable("three", sling::myelin::DT_FLOAT, {1});
constexpr float kThree = 3.0;
three->SetData(&kThree, sizeof(float));
sling::myelin::Flow::Variable *four =
flow_.AddVariable("four", sling::myelin::DT_FLOAT, {1});
constexpr float kFour = 4.0;
four->SetData(&kFour, sizeof(float));
sling::myelin::Flow::Variable *axis =
flow_.AddVariable("axis", sling::myelin::DT_INT32, {1});
constexpr int32 kAxis = 0;
axis->SetData(&kAxis, sizeof(int32));
sling::myelin::Flow::Variable *local_1 =
flow_.AddVariable("local_1", sling::myelin::DT_FLOAT, {3});
sling::myelin::Flow::Variable *local_2 =
flow_.AddVariable("local_2", sling::myelin::DT_FLOAT, {3});
sling::myelin::Flow::Variable *output =
flow_.AddVariable("output", sling::myelin::DT_FLOAT, {6});
output->out = true;
sling::myelin::Flow::Operation *concat_1 = flow_.AddOperation(
function_, "concat_1", "ConcatV2", {one, input, two, axis}, {local_1});
concat_1->SetAttr("N", 3);
sling::myelin::Flow::Operation *concat_2 = flow_.AddOperation(
function_, "concat_2", "ConcatV2", {three, four, input, axis}, {local_2});
concat_2->SetAttr("N", 3);
sling::myelin::Flow::Operation *concat_3 = flow_.AddOperation(
function_, "concat_3", "ConcatV2", {local_1, local_2, axis}, {output});
concat_3->SetAttr("N", 2);
UniqueVector<float> input_feed(1);
(*input_feed)[0] = 0.75;
const std::map<string, MutableAlignedView> feeds = {
{"input", input_feed.view()}};
const CellTrace expected_trace = ParseCellTrace(R"(
name: ')", kCellName, R"('
tensor {
name: 'input'
type: 'float32'
dimension: [1]
aligned_dimension: [1]
order: ORDER_ROW_MAJOR
value: [0.75]
}
tensor {
name: 'local_1'
type: 'float32'
dimension: [3]
aligned_dimension: [3]
order: ORDER_ROW_MAJOR
value: [1.0, 0.75, 2.0]
}
tensor {
name: 'local_2'
type: 'float32'
dimension: [3]
aligned_dimension: [3]
order: ORDER_ROW_MAJOR
value: [3.0, 4.0, 0.75]
}
tensor {
name: 'output'
type: 'float32'
dimension: [6]
aligned_dimension: [6]
order: ORDER_ROW_MAJOR
value: [1.0, 0.75, 2.0, 3.0, 4.0, 0.75]
}
operation {
name: 'concat_1'
type: 'ConcatV2'
kernel: 'BasicConcat'
input: ['one', 'input', 'two', 'axis']
output: ['local_1']
}
operation {
name: 'concat_2'
type: 'ConcatV2'
kernel: 'BasicConcat'
input: ['three', 'four', 'input', 'axis']
output: ['local_2']
}
operation {
name: 'concat_3'
type: 'ConcatV2'
kernel: 'BasicConcat'
input: ['local_1', 'local_2', 'axis']
output: ['output']
}
)");
EXPECT_THAT(GetTrace(feeds), test::EqualsProto(expected_trace));
}
// Tests tracing on a flow that contains an unsupported type: complex128. In
// this case, the tensor values will be missing, but the rest of the trace is
// still extracted.
TEST_F(TraceMyelinInstanceTest, UnsupportedType) {
sling::myelin::Flow::Variable *input =
flow_.AddVariable("input", sling::myelin::DT_COMPLEX128, {1});
input->in = true;
input->ref = true;
sling::myelin::Flow::Variable *zero =
flow_.AddVariable("zero", sling::myelin::DT_COMPLEX128, {1});
const std::vector<char> bytes(2 * sizeof(uint64));
zero->SetData(bytes.data(), bytes.size());
sling::myelin::Flow::Variable *axis =
flow_.AddVariable("axis", sling::myelin::DT_INT32, {1});
constexpr int32 kAxis = 0;
axis->SetData(&kAxis, sizeof(int32));
sling::myelin::Flow::Variable *output =
flow_.AddVariable("output", sling::myelin::DT_COMPLEX128, {2});
output->out = true;
output->ref = true;
sling::myelin::Flow::Operation *concat = flow_.AddOperation(
function_, "concat", "ConcatV2", {input, zero, axis}, {output});
concat->SetAttr("N", 2);
// Both the input and output are refs and need to be fed.
UniqueVector<char> input_feed(2 * sizeof(uint64));
UniqueVector<char> output_feed(4 * sizeof(uint64));
const std::map<string, MutableAlignedView> feeds = {
{"input", input_feed.view()}, //
{"output", output_feed.view()}};
memset(input_feed->data(), 0, input_feed->size());
const CellTrace expected_trace = ParseCellTrace(R"(
name: ')", kCellName, R"('
tensor {
name: 'input'
type: 'complex128'
dimension: [1]
aligned_dimension: [1]
order: ORDER_ROW_MAJOR
}
tensor {
name: 'output'
type: 'complex128'
dimension: [2]
aligned_dimension: [2]
order: ORDER_ROW_MAJOR
}
operation {
name: 'concat'
type: 'ConcatV2'
kernel: 'BasicConcat'
input: ['input', 'zero', 'axis']
output: ['output']
}
)");
EXPECT_THAT(GetTrace(feeds), test::EqualsProto(expected_trace));
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelination.h"
#include <algorithm>
#include <memory>
#include <utility>
#include <vector>
#include "dragnn/protos/spec.pb.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/myelin/myelin_cell_converter.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/trained_model.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Updates the Component subclass in the |component_spec| to a Myelin-based
// version. On error, returns non-OK and modifies nothing.
tensorflow::Status MyelinateComponentSubclass(ComponentSpec *component_spec) {
const string subclass = GetNormalizedComponentBuilderName(*component_spec);
if (subclass != "DynamicComponent") {
return tensorflow::errors::Unimplemented(
"No Myelin-based version of Component subclass '", subclass, "'");
}
// By convention, the Myelin-based version of "FooComponent" should be named
// "MyelinFooComponent".
component_spec->mutable_component_builder()->set_registered_name(
tensorflow::strings::StrCat("Myelin", subclass));
return tensorflow::Status::OK();
}
// Appends the list of component specs in the |master_spec| whose names match
// |component_names| to |matching_components|. On error, returns non-OK.
tensorflow::Status GetMatchingComponentSpecs(
const std::set<string> &component_names, MasterSpec *master_spec,
std::vector<ComponentSpec *> *matching_components) {
// Index the components in the |master_spec| by name.
std::map<string, ComponentSpec *> components;
for (ComponentSpec &component_spec : *master_spec->mutable_component()) {
if (!components.emplace(component_spec.name(), &component_spec).second) {
return tensorflow::errors::InvalidArgument("Duplicate component name: ",
component_spec.name());
}
}
// Append the components named in the |component_names|.
for (const string &component_name : component_names) {
if (components.find(component_name) == components.end()) {
return tensorflow::errors::InvalidArgument("Unknown component name: ",
component_name);
}
matching_components->push_back(components[component_name]);
}
return tensorflow::Status::OK();
}
} // namespace
tensorflow::Status MyelinateCells(const string &saved_model_dir,
const string &master_spec_path,
const std::set<string> &component_names,
const string &output_dir) {
MasterSpec master_spec;
TF_RETURN_IF_ERROR(tensorflow::ReadTextProto(tensorflow::Env::Default(),
master_spec_path, &master_spec));
std::vector<ComponentSpec *> components;
TF_RETURN_IF_ERROR(
GetMatchingComponentSpecs(component_names, &master_spec, &components));
// Returns the path to the output Flow file for the |component_spec|.
const auto get_flow_path = [&](const ComponentSpec &component_spec) {
return tensorflow::io::JoinPath(
output_dir,
tensorflow::strings::StrCat(component_spec.name(), ".flow"));
};
// Modify the MasterSpec first, to catch issues before loading the trained
// model, which is slow.
for (ComponentSpec *component_spec : components) {
// Add a resource for the Flow file to each component. The file will be
// created in a second pass, after loading the trained model.
TF_RETURN_IF_ERROR(
AddMyelinFlowResource(get_flow_path(*component_spec), component_spec));
// Replace the Component subclass with a Myelin-based version.
TF_RETURN_IF_ERROR(MyelinateComponentSubclass(component_spec));
// Set embedding_dim=-1 for all channels.
for (auto &fixed_channel : *component_spec->mutable_fixed_feature()) {
fixed_channel.set_embedding_dim(-1);
}
for (auto &linked_channel : *component_spec->mutable_linked_feature()) {
linked_channel.set_embedding_dim(-1);
}
}
// Write the updated MasterSpec.
TF_RETURN_IF_ERROR(
tensorflow::Env::Default()->RecursivelyCreateDir(output_dir));
TF_RETURN_IF_ERROR(tensorflow::WriteTextProto(
tensorflow::Env::Default(),
tensorflow::io::JoinPath(output_dir, "master-spec"), master_spec));
// Convert each component into a Flow and write it.
TrainedModel trained_model;
TF_RETURN_IF_ERROR(trained_model.Reset(saved_model_dir));
for (const ComponentSpec *component_spec : components) {
string flow_data;
TF_RETURN_IF_ERROR(MyelinCellConverter::Convert(component_spec->name(),
trained_model, &flow_data));
TF_RETURN_IF_ERROR(tensorflow::WriteStringToFile(
tensorflow::Env::Default(), get_flow_path(*component_spec), flow_data));
}
return tensorflow::Status::OK();
}
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
// Utils for modifying pre-trained models to use Myelin.
#ifndef DRAGNN_RUNTIME_MYELIN_MYELINATION_H_
#define DRAGNN_RUNTIME_MYELIN_MYELINATION_H_
#include <set>
#include <string>
#include "syntaxnet/base.h"
#include "tensorflow/core/lib/core/status.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
// Modifies a DRAGNN model to use Myelin.
//
// Loads a TF SavedModel from the |saved_model_dir| and a text-format MasterSpec
// from the |master_spec_path|. Converts each component in |component_names|
// into a Myelin Flow (see myelin_cell_converter.h) and writes the results to
// the |output_dir| as files "<output_dir>/<component_name>.flow". Modifies the
// relevant ComponentSpecs in the MasterSpec to use Myelin as described below,
// and writes it to "<output_dir>/master-spec".
//
// MasterSpec modifications:
// * Adds a resource to each ComponentSpec that points at the relevant Flow file
// in the |output_dir|.
// * Replaces the Component subclass specified in each ComponentSpec with the
// Myelin-based equivalent, which should be named "Myelin<subclass_name>";
// e.g., MyelinDynamicComponent.
// * Sets FixedFeatureChannel.embedding_dim to -1 in all channels, because
// Myelin takes feature IDs as input instead of fixed embedding sums.
// * Sets LinkedFeatureChannel.embedding_dim to -1 in all channels, because
// Myelin handles the linked embedding matrix multiplication (if any) and
// always takes the original activation vector as input.
//
// On error, returns non-OK. Possible errors include:
// * Any file I/O or proto parsing error.
// * The MasterSpec has a duplicate component name.
// * One of the |component_names| does not match anything in the MasterSpec.
// * The MasterSpec already has Myelin Flow resources.
// * One of the components is not supported by Myelin.
// * Error raised by MyelinCellConverter during conversion.
//
// Side note: This function has a file-path-based API so it can be easily
// wrapped in a stand-alone binary.
tensorflow::Status MyelinateCells(const string &saved_model_dir,
const string &master_spec_path,
const std::set<string> &component_names,
const string &output_dir);
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
#endif // DRAGNN_RUNTIME_MYELIN_MYELINATION_H_
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include "dragnn/runtime/myelin/myelination.h"
#include <memory>
#include <string>
#include <utility>
#include "dragnn/core/test/generic.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "syntaxnet/base.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// Arbitrary bogus path.
constexpr char kInvalidPath[] = "path/to/some/invalid/file";
// Relative path to a MasterSpec.
constexpr char kMasterSpecPath[] =
"dragnn/runtime/testdata/rnn_tagger/assets.extra/master_spec";
// Relative path to a saved model.
constexpr char kSavedModelDir[] = "dragnn/runtime/testdata/rnn_tagger";
// Relative path to a directory containing expected output.
constexpr char kExpectedOutputDir[] =
"dragnn/runtime/myelin/testdata/myelination_output";
// Local relative path to the expected output directory.
constexpr char kLocalOutputDir[] =
"dragnn/runtime/myelin/testdata/myelination_output";
// Returns the set of components in the MasterSpec at |kMasterSpecPath|.
std::set<string> GetComponentNames() { return {"rnn", "tagger"}; }
// Returns the path to a test input denoted by the |relative_path|.
string GetInput(const string &relative_path) {
return tensorflow::io::JoinPath(test::GetTestDataPrefix(), relative_path);
}
// Returns a unique output directory for tests.
string GetUniqueOutputDir() {
static int counter = 0;
return tensorflow::io::JoinPath(
tensorflow::testing::TmpDir(),
tensorflow::strings::StrCat("output_", counter++));
}
// Compares the content of the file named |basename| in the |actual_output_dir|
// with the file with the same |basename| in |kExpectedOutputDir|. Can also be
// modified to write the actual file content to |kLocalOutputDir|, for updating
// test expectations.
void CompareOrRewriteTestData(const string &actual_output_dir,
const string &basename) {
string actual_data;
TF_ASSERT_OK(tensorflow::ReadFileToString(
tensorflow::Env::Default(),
tensorflow::io::JoinPath(actual_output_dir, basename), &actual_data));
if (false) {
TF_ASSERT_OK(tensorflow::WriteStringToFile(
tensorflow::Env::Default(),
tensorflow::io::JoinPath(kLocalOutputDir, basename), actual_data));
} else {
string expected_data;
TF_ASSERT_OK(tensorflow::ReadFileToString(
tensorflow::Env::Default(),
GetInput(tensorflow::io::JoinPath(kExpectedOutputDir, basename)),
&expected_data));
// Avoid EXPECT_EQ(), which produces a text diff on error. The diff is not
// interpretable because Flow files are binary, and the test can OOM when it
// tries to diff two large binary files.
EXPECT_TRUE(actual_data == expected_data);
}
}
// Reads a text-format MasterSpec from the |master_spec_path|, clears resource
// file patterns, and writes it back to the |master_spec_path|. The resource
// file patterns would otherwise cause spurious mismatches.
void ClearResourceFilePatterns(const string &master_spec_path) {
MasterSpec master_spec;
TF_ASSERT_OK(tensorflow::ReadTextProto(tensorflow::Env::Default(),
master_spec_path, &master_spec));
for (ComponentSpec &component_spec : *master_spec.mutable_component()) {
for (Resource &resource : *component_spec.mutable_resource()) {
for (Part &part : *resource.mutable_part()) {
part.clear_file_pattern();
}
}
}
TF_ASSERT_OK(tensorflow::WriteTextProto(tensorflow::Env::Default(),
master_spec_path, master_spec));
}
// Tests that MyelinateCells() fails if the saved model is invalid.
TEST(MyelinateCellsTest, InvalidSavedModel) {
EXPECT_FALSE(MyelinateCells(kInvalidPath, GetInput(kMasterSpecPath), {},
GetUniqueOutputDir())
.ok());
}
// Tests that MyelinateCells() fails if the master spec is invalid.
TEST(MyelinateCellsTest, InvalidMasterSpec) {
EXPECT_FALSE(MyelinateCells(GetInput(kSavedModelDir), kInvalidPath, {},
GetUniqueOutputDir())
.ok());
}
// Tests that MyelinateCells() fails if the MasterSpec contains a duplicate
// component.
TEST(MyelinateCellsTest, DuplicateComponent) {
const string kSpec = "component { name:'foo' } component { name:'foo' }";
const string master_spec_path = tensorflow::io::JoinPath(
tensorflow::testing::TmpDir(), "master-spec-with-duplicate");
TF_ASSERT_OK(tensorflow::WriteStringToFile(tensorflow::Env::Default(),
master_spec_path, kSpec));
EXPECT_THAT(MyelinateCells(GetInput(kSavedModelDir), master_spec_path, {},
GetUniqueOutputDir()),
test::IsErrorWithSubstr("Duplicate component name: foo"));
}
// Tests that MyelinateCells() fails if one of the requested components does not
// appear in the MasterSpec.
TEST(MyelinateCellsTest, FilterWithUnknownComponent) {
const string kSpec = "component { name:'foo' } component { name:'bar' }";
const string master_spec_path = tensorflow::io::JoinPath(
tensorflow::testing::TmpDir(), "master-spec-foo-bar");
TF_ASSERT_OK(tensorflow::WriteStringToFile(tensorflow::Env::Default(),
master_spec_path, kSpec));
EXPECT_THAT(MyelinateCells(GetInput(kSavedModelDir), master_spec_path,
{"missing"}, GetUniqueOutputDir()),
test::IsErrorWithSubstr("Unknown component name: missing"));
}
// Tests that MyelinateCells() fails if a component already has a Myelin Flow.
TEST(MyelinateCellsTest, AlreadyHasFlow) {
const string kSpec =
tensorflow::strings::StrCat("component { name: 'foo' resource { name: '",
kMyelinFlowResourceName, "' } }");
const string master_spec_path = tensorflow::io::JoinPath(
tensorflow::testing::TmpDir(), "master-spec-with-flows");
TF_ASSERT_OK(tensorflow::WriteStringToFile(tensorflow::Env::Default(),
master_spec_path, kSpec));
EXPECT_THAT(
MyelinateCells(GetInput(kSavedModelDir), master_spec_path, {"foo"},
GetUniqueOutputDir()),
test::IsErrorWithSubstr("already contains a Myelin Flow resource"));
}
// Tests that MyelinateCells() fails on the wrong Component type.
TEST(MyelinateCellsTest, WrongComponentType) {
const string kSpec =
"component { name: 'foo' component_builder { registered_name: "
"'WrongComponent' } }";
const string master_spec_path =
tensorflow::io::JoinPath(tensorflow::testing::TmpDir(), "master-spec");
TF_ASSERT_OK(tensorflow::WriteStringToFile(tensorflow::Env::Default(),
master_spec_path, kSpec));
EXPECT_THAT(
MyelinateCells(GetInput(kSavedModelDir), master_spec_path, {"foo"},
GetUniqueOutputDir()),
test::IsErrorWithSubstr(
"No Myelin-based version of Component subclass 'WrongComponent'"));
}
// Tests that MyelinateCells() succeeds on the pre-trained inputs and reproduces
// expected outputs.
TEST(MyelinateCellsTest, RegressionTest) {
const string output_dir = GetUniqueOutputDir();
TF_ASSERT_OK(MyelinateCells(GetInput(kSavedModelDir),
GetInput(kMasterSpecPath), GetComponentNames(),
output_dir));
ClearResourceFilePatterns(
tensorflow::io::JoinPath(output_dir, "master-spec"));
CompareOrRewriteTestData(output_dir, "master-spec");
for (const string &component_name : GetComponentNames()) {
const string flow_basename =
tensorflow::strings::StrCat(component_name, ".flow");
CompareOrRewriteTestData(output_dir, flow_basename);
}
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include <stddef.h>
#include <string>
#include <vector>
#include "dragnn/core/compute_session.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/protos/trace.pb.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/extensions.h"
#include "dragnn/runtime/math/types.h"
#include "dragnn/runtime/myelin/myelin_dynamic_component_base.h"
#include "dragnn/runtime/network_states.h"
#include "dragnn/runtime/sequence_features.h"
#include "dragnn/runtime/sequence_links.h"
#include "dragnn/runtime/sequence_model.h"
#include "dragnn/runtime/session_state.h"
#include "dragnn/runtime/variable_store.h"
#include "syntaxnet/base.h"
#include "sling/myelin/compute.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
// A Myelin-based version of DynamicComponent for sequence-based models.
class SequenceMyelinDynamicComponent : public MyelinDynamicComponentBase {
public:
// Implements Component.
tensorflow::Status Initialize(const ComponentSpec &component_spec,
VariableStore *variable_store,
NetworkStateManager *network_state_manager,
ExtensionManager *extension_manager) override;
tensorflow::Status Evaluate(SessionState *session_state,
ComputeSession *compute_session,
ComponentTrace *component_trace) const override;
protected:
// Implements Component.
bool Supports(const ComponentSpec &component_spec,
const string &normalized_builder_name) const override;
bool PreferredTo(const Component &) const override { return false; }
private:
// Binds the fixed feature IDs for the |target_index|'th element of the
// |features| to the |instance|. Uses locals in the |network_states|.
void BindInputIds(const SequenceFeatures &features, int target_index,
const NetworkStates &network_states,
sling::myelin::Instance *instance) const;
// Binds the linked embeddings for the |target_index|'th element in the
// |links| to the |instance|.
void BindInputLinks(const SequenceLinks &links, int target_index,
sling::myelin::Instance *instance) const;
// Sequence-based model evaluator.
SequenceModel sequence_model_;
// Intermediate values used by sequence models.
SharedExtensionHandle<SequenceModel::EvaluateState> evaluate_state_handle_;
};
bool SequenceMyelinDynamicComponent::Supports(
const ComponentSpec &component_spec,
const string &normalized_builder_name) const {
return normalized_builder_name == "SequenceMyelinDynamicComponent" &&
SequenceModel::Supports(component_spec);
}
tensorflow::Status SequenceMyelinDynamicComponent::Initialize(
const ComponentSpec &component_spec, VariableStore *variable_store,
NetworkStateManager *network_state_manager,
ExtensionManager *extension_manager) {
// Initialize the base class first, so its FixedEmbeddingManager and
// LinkedEmbeddingManager can be wrapped in sequence-based versions.
TF_RETURN_IF_ERROR(MyelinDynamicComponentBase::Initialize(
component_spec, variable_store, network_state_manager,
extension_manager));
TF_RETURN_IF_ERROR(sequence_model_.Initialize(
component_spec, kLogitsName, &fixed_embedding_manager(),
&linked_embedding_manager(), network_state_manager));
extension_manager->GetShared(&evaluate_state_handle_);
return tensorflow::Status::OK();
}
void SequenceMyelinDynamicComponent::BindInputIds(
const SequenceFeatures &features, int target_index,
const NetworkStates &network_states,
sling::myelin::Instance *instance) const {
for (size_t channel_id = 0; channel_id < features.num_channels();
++channel_id) {
const MutableVector<int32> id_vector = network_states.GetLocal(
fixed_embedding_manager().id_handle(channel_id, 0));
id_vector[0] = features.GetId(channel_id, target_index);
BindInput(Vector<int32>(id_vector), input_ids()[channel_id].id, instance);
}
}
void SequenceMyelinDynamicComponent::BindInputLinks(
const SequenceLinks &links, int target_index,
sling::myelin::Instance *instance) const {
Vector<float> embedding;
bool is_out_of_bounds = false;
for (size_t channel_id = 0; channel_id < links.num_channels(); ++channel_id) {
links.Get(channel_id, target_index, &embedding, &is_out_of_bounds);
BindInputLink(embedding, is_out_of_bounds, input_links()[channel_id],
instance);
}
}
tensorflow::Status SequenceMyelinDynamicComponent::Evaluate(
SessionState *session_state, ComputeSession *compute_session,
ComponentTrace *component_trace) const {
NetworkStates &network_states = session_state->network_states;
SequenceModel::EvaluateState &state =
session_state->extensions.Get(evaluate_state_handle_);
TF_RETURN_IF_ERROR(
sequence_model_.Preprocess(session_state, compute_session, &state));
// Avoid ComputeSession overhead by directly iterating over the feature IDs.
// Handle forward and reverse iteration via an index and increment.
int target_index = sequence_model_.left_to_right() ? 0 : state.num_steps - 1;
const int target_increment = sequence_model_.left_to_right() ? 1 : -1;
sling::myelin::Instance &instance = GetInstance(session_state);
for (size_t step_index = 0; step_index < state.num_steps;
++step_index, target_index += target_increment) {
// Bind inputs and outputs into the |instance|.
BindInputIds(state.features, target_index, network_states, &instance);
BindInputLinks(state.links, target_index, &instance);
BindInputRecurrences(step_index, network_states, &instance);
BindOutputLayers(step_index, network_states, &instance);
// Invoke the cell in the |instance|.
instance.Compute();
MaybeTrace(step_index, &instance, component_trace);
}
return sequence_model_.Predict(network_states, &state);
}
DRAGNN_RUNTIME_REGISTER_COMPONENT(SequenceMyelinDynamicComponent);
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
#include <memory>
#include <string>
#include <vector>
#include "dragnn/core/compute_session.h"
#include "dragnn/core/input_batch_cache.h"
#include "dragnn/core/test/generic.h"
#include "dragnn/protos/spec.pb.h"
#include "dragnn/protos/trace.pb.h"
#include "dragnn/runtime/component.h"
#include "dragnn/runtime/extensions.h"
#include "dragnn/runtime/math/types.h"
#include "dragnn/runtime/myelin/myelin_spec_utils.h"
#include "dragnn/runtime/network_states.h"
#include "dragnn/runtime/sequence_backend.h"
#include "dragnn/runtime/sequence_extractor.h"
#include "dragnn/runtime/sequence_linker.h"
#include "dragnn/runtime/sequence_predictor.h"
#include "dragnn/runtime/session_state.h"
#include "dragnn/runtime/test/network_test_base.h"
#include "dragnn/runtime/variable_store.h"
#include "syntaxnet/base.h"
#include <gmock/gmock.h>
#include "sling/file/file.h"
#include "sling/myelin/flow.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
namespace syntaxnet {
namespace dragnn {
namespace runtime {
namespace {
using ::testing::Return;
constexpr int kFlowVersion = 4;
constexpr int kNumSteps = 50;
constexpr int kVocabularySize = 123;
constexpr int kFixedDim = 6;
constexpr int kLinkedDim = 4;
constexpr int kLogitsDim = kFixedDim + kLinkedDim;
constexpr char kLogitsName[] = "logits";
constexpr char kPreviousComponentName[] = "previous_component";
constexpr char kPreviousLayerName[] = "previous_layer";
constexpr float kPreviousLayerValue = -1.0;
// Builds and writes a simple Flow file with a function named |function_name|
// that gathers the rows of a matrix, concatenates that with a linked embedding,
// and outputs the result as the classification logits. Each row is filled with
// its index, so we can infer which indices were gathered.
string WriteFlowFile(const string &function_name) {
sling::myelin::Flow flow;
// A fixed feature ID input.
sling::myelin::Flow::Variable *id =
flow.AddVariable("id", sling::myelin::DT_INT32, {1});
id->ref = true;
id->aliases.push_back(MakeMyelinInputFixedFeatureIdName(0, 0));
// A linked feature embedding input.
sling::myelin::Flow::Variable *link =
flow.AddVariable("link", sling::myelin::DT_FLOAT, {1, kLinkedDim});
link->ref = true;
link->aliases.push_back(MakeMyelinInputLinkedActivationVectorName(0));
// An embedding matrix constant. Each embedding is filled with its index.
sling::myelin::Flow::Variable *embeddings = flow.AddVariable(
"embeddings", sling::myelin::DT_FLOAT, {kVocabularySize, kFixedDim});
std::vector<float> data(kVocabularySize * kLogitsDim);
for (int row = 0; row < kVocabularySize; ++row) {
for (int column = 0; column < kFixedDim; ++column) {
data[row * kFixedDim + column] = row;
}
}
embeddings->SetData(data.data(), data.size() * sizeof(float));
// The retrieved embedding row.
sling::myelin::Flow::Variable *row =
flow.AddVariable("row", sling::myelin::DT_FLOAT, {1, kFixedDim});
// A concatenation axis constant.
sling::myelin::Flow::Variable *axis =
flow.AddVariable("axis", sling::myelin::DT_INT32, {1});
const int32 axis_value = 1;
axis->SetData(&axis_value, sizeof(int32));
// The classification logits output.
sling::myelin::Flow::Variable *logits =
flow.AddVariable(kLogitsName, sling::myelin::DT_FLOAT, {1, kLogitsDim});
logits->ref = true;
logits->aliases.push_back(MakeMyelinOutputLayerName(kLogitsName));
// Function that contains the ops and variables.
sling::myelin::Flow::Function *function = flow.AddFunction(function_name);
// A Gather op that looks up the |id| in the |embeddings|, and returns the
// result in the |row|.
flow.AddOperation(function, "gather", "Gather", {embeddings, id}, {row});
// A Concat op that concatenates the |row| and |link| along the |axis|,
// placing the result in the |logits| output.
flow.AddOperation(function, "concat", "ConcatV2", {row, link, axis},
{logits});
const string flow_path =
tensorflow::io::JoinPath(tensorflow::testing::TmpDir(), "foo.flow");
sling::File::Init();
flow.Save(flow_path, kFlowVersion);
return flow_path;
}
// Sequence extractor that extracts [0, 2, 4, ...].
class EvenNumbers : public SequenceExtractor {
public:
// Implements SequenceExtractor.
bool Supports(const FixedFeatureChannel &,
const ComponentSpec &) const override {
return true;
}
tensorflow::Status Initialize(const FixedFeatureChannel &,
const ComponentSpec &) override {
return tensorflow::Status::OK();
}
tensorflow::Status GetIds(InputBatchCache *,
std::vector<int32> *ids) const override {
ids->clear();
for (int i = 0; i < num_steps_; ++i) ids->push_back(2 * i);
return tensorflow::Status::OK();
}
// Sets the number of steps to emit.
static void SetNumSteps(int num_steps) { num_steps_ = num_steps; }
private:
// The number of steps to produce.
static int num_steps_;
};
int EvenNumbers::num_steps_ = kNumSteps;
DRAGNN_RUNTIME_REGISTER_SEQUENCE_EXTRACTOR(EvenNumbers);
// Component that supports a particular component name and is not preferred.
// Used to exercise PreferredTo().
class NotPreferred : public Component {
public:
// Implements Component.
tensorflow::Status Initialize(const ComponentSpec &, VariableStore *,
NetworkStateManager *,
ExtensionManager *) override {
return tensorflow::Status::OK();
}
tensorflow::Status Evaluate(SessionState *, ComputeSession *,
ComponentTrace *) const override {
return tensorflow::Status::OK();
}
bool Supports(const ComponentSpec &spec, const string &) const override {
return spec.name() == "InSupportsConflictTest";
}
bool PreferredTo(const Component &) const override { return false; }
};
DRAGNN_RUNTIME_REGISTER_COMPONENT(NotPreferred);
// Trivial linker that links everything to step 0.
class LinkToZero : public SequenceLinker {
public:
// Implements SequenceLinker.
bool Supports(const LinkedFeatureChannel &,
const ComponentSpec &) const override {
return true;
}
tensorflow::Status Initialize(const LinkedFeatureChannel &,
const ComponentSpec &) override {
return tensorflow::Status::OK();
}
tensorflow::Status GetLinks(size_t, InputBatchCache *,
std::vector<int32> *links) const override {
links->assign(num_steps_, 0);
return tensorflow::Status::OK();
}
// Sets the number of steps to emit.
static void SetNumSteps(int num_steps) { num_steps_ = num_steps; }
private:
// The number of steps to produce.
static int num_steps_;
};
int LinkToZero::num_steps_ = kNumSteps;
DRAGNN_RUNTIME_REGISTER_SEQUENCE_LINKER(LinkToZero);
// Trivial predictor that captures the prediction logits.
class CaptureLogits : public SequencePredictor {
public:
// Implements SequenceLinker.
bool Supports(const ComponentSpec &) const override { return true; }
tensorflow::Status Initialize(const ComponentSpec &) override {
return tensorflow::Status::OK();
}
tensorflow::Status Predict(Matrix<float> logits,
InputBatchCache *) const override {
GetLogits() = logits;
return tensorflow::Status::OK();
}
// Returns the captured logits.
static Matrix<float> &GetLogits() {
static auto *logits = new Matrix<float>();
return *logits;
}
};
DRAGNN_RUNTIME_REGISTER_SEQUENCE_PREDICTOR(CaptureLogits);
class SequenceMyelinDynamicComponentTest : public NetworkTestBase {
protected:
// Adds default call expectations. Since these are added first, they can be
// overridden by call expectations in individual tests.
SequenceMyelinDynamicComponentTest() {
EXPECT_CALL(compute_session_, GetInputBatchCache())
.WillRepeatedly(Return(&input_));
EXPECT_CALL(compute_session_, GetReadiedComponent(kTestComponentName))
.WillRepeatedly(Return(&backend_));
TF_CHECK_OK(Component::CreateOrError("SequenceMyelinDynamicComponent",
&component_));
// Some tests overwrite these; ensure that they are restored to the normal
// values at the start of each test.
EvenNumbers::SetNumSteps(kNumSteps);
LinkToZero::SetNumSteps(kNumSteps);
CaptureLogits::GetLogits() = Matrix<float>();
}
// Build and write the flow file once.
static void SetUpTestCase() {
flow_path_ = new string(WriteFlowFile(kTestComponentName));
}
// Cleans up the flow file path.
static void TearDownTestCase() {
delete flow_path_;
flow_path_ = nullptr;
}
// Creates a component, initializes it based on the |component_spec|, and
// evaluates it. On error, returns non-OK.
tensorflow::Status Run(ComponentSpec component_spec) {
component_spec.set_name(kTestComponentName);
TF_RETURN_IF_ERROR(AddMyelinFlowResource(*flow_path_, &component_spec));
AddComponent(kPreviousComponentName);
AddLayer(kPreviousLayerName, kLinkedDim);
AddComponent(kTestComponentName);
TF_RETURN_IF_ERROR(component_->Initialize(component_spec, &variable_store_,
&network_state_manager_,
&extension_manager_));
network_states_.Reset(&network_state_manager_);
session_state_.extensions.Reset(&extension_manager_);
StartComponent(kNumSteps);
FillLayer(kPreviousComponentName, kPreviousLayerName, kPreviousLayerValue);
StartComponent(0);
TF_RETURN_IF_ERROR(
component_->Evaluate(&session_state_, &compute_session_, nullptr));
return tensorflow::Status::OK();
}
// Returns the sequence size passed to the |backend_|.
int GetBackendSequenceSize() {
// The sequence size is not directly exposed, but can be inferred using one
// of the reverse step translators.
return backend_.GetStepLookupFunction("reverse-token")(0, 0, 0) + 1;
}
// Path to a simple Myelin Flow file.
static const string *flow_path_;
// Component used in the test.
std::unique_ptr<Component> component_;
// Input batch injected into Evaluate() by default.
InputBatchCache input_;
// Backend injected into Evaluate().
SequenceBackend backend_;
};
const string *SequenceMyelinDynamicComponentTest::flow_path_ = nullptr;
// Returns a ComponentSpec that is supported.
ComponentSpec MakeSupportedSpec() {
ComponentSpec component_spec;
component_spec.set_num_actions(kLogitsDim);
component_spec.mutable_component_builder()->set_registered_name(
"SequenceMyelinDynamicComponent");
component_spec.mutable_component_builder()->mutable_parameters()->insert(
{"sequence_extractors", "EvenNumbers"});
component_spec.mutable_component_builder()->mutable_parameters()->insert(
{"sequence_linkers", "LinkToZero"});
component_spec.mutable_component_builder()->mutable_parameters()->insert(
{"sequence_predictor", "CaptureLogits"});
component_spec.mutable_backend()->set_registered_name("SequenceBackend");
FixedFeatureChannel *fixed_feature = component_spec.add_fixed_feature();
fixed_feature->set_size(1);
fixed_feature->set_embedding_dim(-1);
LinkedFeatureChannel *linked_feature = component_spec.add_linked_feature();
linked_feature->set_source_component(kPreviousComponentName);
linked_feature->set_source_layer(kPreviousLayerName);
linked_feature->set_size(1);
linked_feature->set_embedding_dim(-1);
return component_spec;
}
// Tests that the component supports a supported spec.
TEST_F(SequenceMyelinDynamicComponentTest, Supported) {
string component_type;
const ComponentSpec component_spec = MakeSupportedSpec();
TF_ASSERT_OK(Component::Select(component_spec, &component_type));
}
// Tests that the component does not support a spec with the wrong component
// builder.
TEST_F(SequenceMyelinDynamicComponentTest, UnsupportedComponentBuilder) {
string component_type;
ComponentSpec component_spec = MakeSupportedSpec();
component_spec.mutable_component_builder()->set_registered_name("bad");
EXPECT_THAT(Component::Select(component_spec, &component_type),
test::IsErrorWithSubstr("Could not find a best"));
}
// Tests that the component
TEST_F(SequenceMyelinDynamicComponentTest, SupportsConflict) {
string component_type;
ComponentSpec component_spec = MakeSupportedSpec();
component_spec.set_name("InSupportsConflictTest"); // see NotPreferred
EXPECT_THAT(
Component::Select(component_spec, &component_type),
test::IsErrorWithSubstr("both think they should be dis-preferred"));
}
// Asserts that the vector starts with |kFixedDim| copies of |value| and ends
// with |kLinkedDim| copies of |kPreviousLayerValue|.
void AssertOutputRow(Vector<float> row, float value) {
ASSERT_EQ(row.size(), kLogitsDim);
for (int i = 0; i < row.size(); ++i) {
if (i < kFixedDim) {
ASSERT_EQ(row[i], value);
} else {
ASSERT_EQ(row[i], kPreviousLayerValue);
}
}
}
// Tests that the component extracts a left-to-right sequence by default.
TEST_F(SequenceMyelinDynamicComponentTest, LeftToRightByDefault) {
TF_ASSERT_OK(Run(MakeSupportedSpec()));
EXPECT_EQ(GetBackendSequenceSize(), kNumSteps);
const Matrix<float> logits = CaptureLogits::GetLogits();
ASSERT_EQ(logits.num_rows(), kNumSteps);
ASSERT_EQ(logits.num_columns(), kLogitsDim);
for (int i = 0; i < kNumSteps; ++i) {
AssertOutputRow(logits.row(i), 2.0 * i);
}
}
// Tests that the component can be explicitly configured for a left-to-right
// sequence.
TEST_F(SequenceMyelinDynamicComponentTest, LeftToRightExplicitly) {
ComponentSpec component_spec = MakeSupportedSpec();
(*component_spec.mutable_transition_system()
->mutable_parameters())["left_to_right"] = "true";
TF_ASSERT_OK(Run(component_spec));
EXPECT_EQ(GetBackendSequenceSize(), kNumSteps);
const Matrix<float> logits = CaptureLogits::GetLogits();
ASSERT_EQ(logits.num_rows(), kNumSteps);
ASSERT_EQ(logits.num_columns(), kLogitsDim);
for (int i = 0; i < kNumSteps; ++i) {
AssertOutputRow(logits.row(i), 2.0 * i);
}
}
// Tests that the component can be explicitly configured for a right-to-left
// sequence.
TEST_F(SequenceMyelinDynamicComponentTest, RightToLeft) {
ComponentSpec component_spec = MakeSupportedSpec();
(*component_spec.mutable_transition_system()
->mutable_parameters())["left_to_right"] = "false";
TF_ASSERT_OK(Run(component_spec));
EXPECT_EQ(GetBackendSequenceSize(), kNumSteps);
const Matrix<float> logits = CaptureLogits::GetLogits();
ASSERT_EQ(logits.num_rows(), kNumSteps);
ASSERT_EQ(logits.num_columns(), kLogitsDim);
for (int i = 0; i < kNumSteps; ++i) {
const int reversed = kNumSteps - i - 1;
AssertOutputRow(logits.row(i), 2.0 * reversed);
}
}
// Tests that the component can handle an empty sequence.
TEST_F(SequenceMyelinDynamicComponentTest, EmptySequence) {
EvenNumbers::SetNumSteps(0);
LinkToZero::SetNumSteps(0);
TF_ASSERT_OK(Run(MakeSupportedSpec()));
EXPECT_EQ(GetBackendSequenceSize(), 0);
const Matrix<float> logits = CaptureLogits::GetLogits();
ASSERT_EQ(logits.num_rows(), 0);
}
} // namespace
} // namespace runtime
} // namespace dragnn
} // namespace syntaxnet
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment