Count total number of variables.
(reader, count_exclude_pattern="")
| 31 | |
| 32 | |
| 33 | def _count_total_params(reader, count_exclude_pattern=""): |
| 34 | """Count total number of variables.""" |
| 35 | var_to_shape_map = reader.get_variable_to_shape_map() |
| 36 | |
| 37 | # Filter out tensors that we don't want to count |
| 38 | if count_exclude_pattern: |
| 39 | regex_pattern = re.compile(count_exclude_pattern) |
| 40 | new_var_to_shape_map = {} |
| 41 | exclude_num_tensors = 0 |
| 42 | exclude_num_params = 0 |
| 43 | for v in var_to_shape_map: |
| 44 | if regex_pattern.search(v): |
| 45 | exclude_num_tensors += 1 |
| 46 | exclude_num_params += np.prod(var_to_shape_map[v]) |
| 47 | else: |
| 48 | new_var_to_shape_map[v] = var_to_shape_map[v] |
| 49 | var_to_shape_map = new_var_to_shape_map |
| 50 | print("# Excluding %d tensors (%d params) that match %s when counting." % ( |
| 51 | exclude_num_tensors, exclude_num_params, count_exclude_pattern)) |
| 52 | |
| 53 | var_sizes = [np.prod(var_to_shape_map[v]) for v in var_to_shape_map] |
| 54 | return np.sum(var_sizes, dtype=int) |
| 55 | |
| 56 | |
| 57 | def print_tensors_in_checkpoint_file(file_name, tensor_name, all_tensors, |
no test coverage detected