experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add extra VLOGs to save/restore ops.

Browse Source 
PiperOrigin-RevId: 319277619
Change-Id: Ib49c674a6e1a64a16fb96b148721c948be07a8ba
This commit is contained in:
Bruce Fontaine 2020-07-01 12:40:49 -07:00 committed by TensorFlower Gardener
parent 88be89ea85
commit 9214e6dac3
2 changed files with 39 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tensorflow/core/kernels/save_restore_tensor.cc
View File

@ -301,6 +301,24 @@ struct RestoreOp {
TF_RETURN_IF_ERROR(
reader->LookupSlice(tensor_name, parsed_slice, restored_tensor));
}
if (VLOG_IS_ON(5)) {
if (restored_tensor->dtype() == DT_FLOAT) {
const float* t_data = restored_tensor->flat<float>().data();
float min = std::numeric_limits<float>::infinity();
float max = -std::numeric_limits<float>::infinity();
float avg = 0.0;
for (int i = 0; i < restored_tensor->NumElements(); ++i) {
if (t_data[i] < min) min = t_data[i];
if (t_data[i] > max) max = t_data[i];
avg += t_data[i];
}
VLOG(5) << " min " << min << " max " << max << " avg "
<< avg / restored_tensor->NumElements() << " total elts "
<< restored_tensor->NumElements();
}
}
VLOG(1) << "Done restoring tensor " << idx << " : " << tensor_name << " : "
<< restored_full_shape.num_elements();
return Status::OK();
}

21
tensorflow/core/kernels/save_restore_v2_ops.cc
View File

@ -112,6 +112,7 @@ class SaveV2 : public OpKernel {
for (int i = 0; i < num_tensors; ++i) {
const string& tensor_name = tensor_names_flat(i);
const Tensor& tensor = context->input(i + kFixedInputs);
VLOG(2) << "Starting save of " << tensor_name;
if (!shape_and_slices_flat(i).empty()) {
const string& shape_spec = shape_and_slices_flat(i);
@ -133,8 +134,28 @@ class SaveV2 : public OpKernel {
} else {
OP_REQUIRES_OK(context, writer.Add(tensor_name, tensor));
}
if (VLOG_IS_ON(5)) {
if (tensor.dtype() == DT_FLOAT) {
const float* t_data = tensor.flat<float>().data();
float min = std::numeric_limits<float>::infinity();
float max = -std::numeric_limits<float>::infinity();
float avg = 0.0;
for (int i = 0; i < tensor.NumElements(); ++i) {
if (t_data[i] < min) min = t_data[i];
if (t_data[i] > max) max = t_data[i];
avg += t_data[i];
}
VLOG(5) << " min " << min << " max " << max << " avg "
<< avg / tensor.NumElements() << " total elts "
<< tensor.NumElements();
}
}
VLOG(2) << "Done save of " << tensor_name;
}
OP_REQUIRES_OK(context, writer.Finish());
VLOG(1) << "Done BundleWriter, prefix_string: " << prefix_string;
}
};
REGISTER_KERNEL_BUILDER(Name("SaveV2").Device(DEVICE_CPU), SaveV2);