| 97 | |
| 98 | template <typename T> |
| 99 | Status TensorSliceWriter::Add(const string& name, const TensorShape& shape, |
| 100 | const TensorSlice& slice, const T* data) { |
| 101 | // The tensor and the slice have to be compatible |
| 102 | if (shape.dims() != slice.dims()) { |
| 103 | return errors::Internal("Incompatible tensor shape and slice: ", "shape = ", |
| 104 | shape.DebugString(), |
| 105 | ", slice = ", slice.DebugString()); |
| 106 | } |
| 107 | DataType dt = DataTypeToEnum<T>::value; |
| 108 | // We need to add an entry for "name" if there isn't an entry already. |
| 109 | int index = gtl::FindWithDefault(name_to_index_, name, -1); |
| 110 | if (index >= 0) { |
| 111 | // The same tensor has been registered -- we verify that the shapes and the |
| 112 | // type agree. |
| 113 | const SavedSliceMeta& ssm = sts_.meta().tensor(index); |
| 114 | CHECK_EQ(name, ssm.name()) << ProtoShortDebugString(ssm); |
| 115 | TensorShape ssm_shape(ssm.shape()); |
| 116 | if (!shape.IsSameSize(ssm_shape)) { |
| 117 | return errors::Internal( |
| 118 | "Mismatching shapes: existing tensor = ", ssm_shape.DebugString(), |
| 119 | ", trying to add name ", name, ", shape = ", shape.DebugString()); |
| 120 | } |
| 121 | if (dt != ssm.type()) { |
| 122 | return errors::Internal( |
| 123 | "Mismatching types: existing type = ", DataTypeString(ssm.type()), |
| 124 | ", trying to add name ", name, ", type = ", DataTypeString(dt)); |
| 125 | } |
| 126 | } else { |
| 127 | // Insert the new tensor name with the shape information |
| 128 | index = sts_.meta().tensor_size(); |
| 129 | name_to_index_.insert(std::make_pair(name, index)); |
| 130 | SavedSliceMeta* ssm = sts_.mutable_meta()->add_tensor(); |
| 131 | ssm->set_name(name); |
| 132 | shape.AsProto(ssm->mutable_shape()); |
| 133 | ssm->set_type(dt); |
| 134 | } |
| 135 | // Now we need to add the slice info the list of slices. |
| 136 | SavedSliceMeta* ssm = sts_.mutable_meta()->mutable_tensor(index); |
| 137 | slice.AsProto(ssm->add_slice()); |
| 138 | |
| 139 | // Now we need to add the real data. |
| 140 | { |
| 141 | SavedTensorSlices sts; |
| 142 | SavedSlice* ss = sts.mutable_data(); |
| 143 | ss->set_name(name); |
| 144 | slice.AsProto(ss->mutable_slice()); |
| 145 | TensorShape saved_shape(ssm->shape()); |
| 146 | TensorShape sliced_shape; |
| 147 | TF_RETURN_IF_ERROR(slice.SliceTensorShape(saved_shape, &sliced_shape)); |
| 148 | TF_RETURN_IF_ERROR(SaveData(data, sliced_shape.num_elements(), ss)); |
| 149 | string key = EncodeTensorNameSlice(name, slice); |
| 150 | // TODO(yangke): consider doing a two-pass thing where the first pass just |
| 151 | // list the tensor slices we want to save and then another pass to actually |
| 152 | // set the data. Need to figure out if the interface works well. |
| 153 | std::pair<string, string> key_value(key, ""); |
| 154 | if (!sts.AppendToString(&key_value.second)) { |
| 155 | return errors::Internal("Error writing Tensor. Possible size overflow."); |
| 156 | } |