Writes weights to a binary format on disk for ingestion by JavaScript. Weights are organized into groups. When writing to disk, the bytes from all weights in each group are concatenated together and then split into shards (default is 4MB). This means that large weights (> shard_size) ge
(
weight_groups, write_dir, shard_size_bytes=1024 * 1024 * 4,
write_manifest=True, quantization_dtype_map=None)
| 33 | np.dtype(np.complex128): np.complex64} |
| 34 | |
| 35 | def write_weights( |
| 36 | weight_groups, write_dir, shard_size_bytes=1024 * 1024 * 4, |
| 37 | write_manifest=True, quantization_dtype_map=None): |
| 38 | """Writes weights to a binary format on disk for ingestion by JavaScript. |
| 39 | |
| 40 | Weights are organized into groups. When writing to disk, the bytes from all |
| 41 | weights in each group are concatenated together and then split into shards |
| 42 | (default is 4MB). This means that large weights (> shard_size) get sharded |
| 43 | and small weights (< shard_size) will be packed. If the bytes can't be split |
| 44 | evenly into shards, there will be a leftover shard that is smaller than the |
| 45 | shard size. |
| 46 | |
| 47 | Weights are optionally quantized to either 8 or 16 bits for compression, |
| 48 | which is enabled via the `quantization_dtype_map`. |
| 49 | |
| 50 | Args: |
| 51 | weight_groups: An list of groups. Each group is an array of weight |
| 52 | entries. Each entry is a dict that maps a unique name to a numpy array, |
| 53 | for example: |
| 54 | entry = { |
| 55 | 'name': 'weight1', |
| 56 | 'data': np.array([1, 2, 3], 'float32') |
| 57 | } |
| 58 | |
| 59 | Weights groups would then look like: |
| 60 | weight_groups = [ |
| 61 | [group_0_entry1, group_0_entry2], |
| 62 | [group_1_entry1, group_1_entry2], |
| 63 | ] |
| 64 | |
| 65 | The 'name' must be unique across all groups and all entries. The 'data' |
| 66 | field must be a numpy ndarray. |
| 67 | write_dir: A directory to write the files to. |
| 68 | shard_size_bytes: The size of shards in bytes. Defaults to 4MB, which is |
| 69 | the max file size for caching for all major browsers. |
| 70 | write_manifest: Whether to write the manifest JSON to disk. Defaults to |
| 71 | True. |
| 72 | quantization_dtype_map: (Optional) A mapping from dtype |
| 73 | (`uint8`, `uint16`, `float16`) to weights names. The weight mapping |
| 74 | supports wildcard substitution. |
| 75 | Returns: |
| 76 | The weights manifest JSON dict. |
| 77 | |
| 78 | An example manifest with 2 groups, 2 weights, and each weight sharded |
| 79 | into 2: |
| 80 | |
| 81 | The manifest JSON looks like the following: |
| 82 | [{ |
| 83 | 'paths': ['group1-shard1of2', 'group1-shard2of2'], |
| 84 | 'weights': [{ |
| 85 | 'name': 'weight1', |
| 86 | 'shape': [1000, 1000], |
| 87 | 'dtype': 'float32' |
| 88 | }] |
| 89 | }, { |
| 90 | 'paths': ['group2-shard1of2', 'group2-shard2of2'], |
| 91 | 'weights': [{ |
| 92 | 'name': 'weight2', |
nothing calls this directly
no test coverage detected
searching dependent graphs…