| 143 | |
| 144 | template <typename T> |
| 145 | bool TensorSliceReader::CopySliceData(const string& name, |
| 146 | const TensorSlice& slice, T* data) const { |
| 147 | std::vector<std::pair<TensorSlice, string>> details; |
| 148 | const TensorSliceSet* tss; |
| 149 | { |
| 150 | mutex_lock l(mu_); |
| 151 | tss = FindTensorSlice(name, slice, &details); |
| 152 | if (!tss && !all_shards_loaded_) { |
| 153 | VLOG(1) << "Did not find slice in preferred shard, loading all shards." |
| 154 | << name << ": " << slice.DebugString(); |
| 155 | LoadAllShards(); |
| 156 | tss = FindTensorSlice(name, slice, &details); |
| 157 | } |
| 158 | if (!tss) { |
| 159 | // No such tensor |
| 160 | return false; |
| 161 | } |
| 162 | } |
| 163 | // We have the data -- copy it over. |
| 164 | string value; |
| 165 | for (const auto& x : details) { |
| 166 | const TensorSlice& slice_s = x.first; |
| 167 | const string& fname = x.second; |
| 168 | int idx = gtl::FindWithDefault(fname_to_index_, fname, -1); |
| 169 | CHECK_GE(idx, 0) << "Failed to find the index for filename " << fname; |
| 170 | // We read a record in the corresponding sstable |
| 171 | const string key = EncodeTensorNameSlice(name, slice_s); |
| 172 | if (!sss_[idx]->Get(key, &value)) { |
| 173 | VLOG(1) << "Failed to seek to the record for tensor " << name |
| 174 | << ", slice " << slice_s.DebugString() |
| 175 | << ": computed key = " << key; |
| 176 | return false; |
| 177 | } |
| 178 | SavedTensorSlices sts; |
| 179 | if (!ParseProtoUnlimited(&sts, value)) { |
| 180 | VLOG(1) << "Failed to parse the record for tensor " << name << ", slice " |
| 181 | << slice_s.DebugString() << ": computed key = " << key; |
| 182 | return false; |
| 183 | } |
| 184 | // Ensure the TensorSlice contains the expected amount of data. |
| 185 | TensorShape shp_s; |
| 186 | Status s = slice_s.SliceTensorShape(tss->shape(), &shp_s); |
| 187 | if (!s.ok()) { |
| 188 | VLOG(1) << "Failed to slice tensor " << name << ", slice " |
| 189 | << slice_s.DebugString() << ": " << s; |
| 190 | return false; |
| 191 | } |
| 192 | if (checkpoint::TensorProtoDataSize<T>(sts.data().data()) != |
| 193 | shp_s.num_elements()) { |
| 194 | VLOG(1) << "Tensor " << name << ", slice " << slice_s.DebugString() |
| 195 | << " had an unexpected amount of data: expected = " |
| 196 | << shp_s.num_elements() << ", got = " |
| 197 | << checkpoint::TensorProtoDataSize<T>(sts.data().data()); |
| 198 | return false; |
| 199 | } |
| 200 | CopyDataFromTensorSliceToTensorSlice( |
| 201 | tss->shape(), slice_s, slice, |
| 202 | checkpoint::TensorProtoData<T>(sts.data().data()), data); |