Print a description of the current model parameters. Skip variables starting with "tower", as they are just duplicates built by data-parallel logic.
()
| 13 | |
| 14 | |
| 15 | def describe_trainable_vars(): |
| 16 | """ |
| 17 | Print a description of the current model parameters. |
| 18 | Skip variables starting with "tower", as they are just duplicates built by data-parallel logic. |
| 19 | """ |
| 20 | train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) |
| 21 | if len(train_vars) == 0: |
| 22 | logger.warn("No trainable variables in the graph!") |
| 23 | return |
| 24 | total = 0 |
| 25 | total_bytes = 0 |
| 26 | data = [] |
| 27 | for v in train_vars: |
| 28 | if v.name.startswith('tower'): |
| 29 | continue |
| 30 | shape = v.get_shape() |
| 31 | ele = shape.num_elements() |
| 32 | if ele is None: |
| 33 | logger.warn("Shape of variable {} is not fully defined but {}.".format(v.name, shape)) |
| 34 | ele = 0 |
| 35 | try: |
| 36 | shape = shape.as_list() |
| 37 | except ValueError: |
| 38 | shape = '<unknown>' |
| 39 | |
| 40 | total += ele |
| 41 | total_bytes += ele * v.dtype.size |
| 42 | data.append([get_op_tensor_name(v.name)[0], shape, ele, v.device, v.dtype.base_dtype.name]) |
| 43 | headers = ['name', 'shape', '#elements', 'device', 'dtype'] |
| 44 | |
| 45 | dtypes = list({x[4] for x in data}) |
| 46 | if len(dtypes) == 1 and dtypes[0] == "float32": |
| 47 | # don't log the dtype if all vars are float32 (default dtype) |
| 48 | for x in data: |
| 49 | del x[4] |
| 50 | del headers[4] |
| 51 | |
| 52 | devices = {x[3] for x in data} |
| 53 | if len(devices) == 1: |
| 54 | # don't log the device if all vars on the same device |
| 55 | for x in data: |
| 56 | del x[3] |
| 57 | del headers[3] |
| 58 | |
| 59 | table = tabulate(data, headers=headers) |
| 60 | |
| 61 | size_mb = total_bytes / 1024.0**2 |
| 62 | summary_msg = colored( |
| 63 | "\nNumber of trainable variables: {}".format(len(data)) + |
| 64 | "\nNumber of parameters (elements): {}".format(total) + |
| 65 | "\nStorage space needed for all trainable variables: {:.02f}MB".format(size_mb), |
| 66 | 'cyan') |
| 67 | logger.info(colored("List of Trainable Variables: \n", 'cyan') + table + summary_msg) |
| 68 | |
| 69 | |
| 70 | def get_shape_str(tensors): |
no test coverage detected