Given a list of dict, each of which contains torch.Tensor or a list of torch.Tensor, we concat along `dim` and create a dict by concat each of them. Args: dict_list: dim_dict: Returns: a dict having the same keys
(
dict_list: T.List[T.Dict[str, T.Union[torch.Tensor, T.List[torch.Tensor]]]],
dim_dict: T.Union[int, T.Dict[str, int]],
)
| 114 | |
| 115 | |
| 116 | def cat_dict( |
| 117 | dict_list: T.List[T.Dict[str, T.Union[torch.Tensor, T.List[torch.Tensor]]]], |
| 118 | dim_dict: T.Union[int, T.Dict[str, int]], |
| 119 | ) -> T.Dict[str, torch.Tensor]: |
| 120 | """ |
| 121 | Given a list of dict, each of which contains torch.Tensor or a list of torch.Tensor, |
| 122 | we concat along `dim` and create a dict by concat each of them. |
| 123 | |
| 124 | Args: |
| 125 | dict_list: |
| 126 | dim_dict: |
| 127 | |
| 128 | Returns: |
| 129 | a dict having the same keys |
| 130 | """ |
| 131 | if len(dict_list) == 0: |
| 132 | return dict() |
| 133 | |
| 134 | if isinstance(dim_dict, int): |
| 135 | dim = dim_dict |
| 136 | dim_dict = dict() |
| 137 | for key in dict_list[0]: |
| 138 | dim_dict[key] = dim |
| 139 | |
| 140 | out_dict = dict() |
| 141 | for key in dict_list[0]: |
| 142 | out_dict[key] = [d[key] for d in dict_list] |
| 143 | |
| 144 | for key in out_dict: |
| 145 | if out_dict[key][0] is None: |
| 146 | out_dict[key] = None |
| 147 | elif isinstance(out_dict[key][0], torch.Tensor): |
| 148 | out_dict[key] = torch.cat(out_dict[key], dim=dim_dict[key]) |
| 149 | else: # out_dict[key] is a list of "list of tensor" |
| 150 | batch_num = len(out_dict[key]) |
| 151 | tensor_num = len(out_dict[key][0]) |
| 152 | feature_list = [''] * tensor_num # initialize list |
| 153 | for tensor_id in range(tensor_num): |
| 154 | feature_list[tensor_id] = torch.cat( |
| 155 | [out_dict[key][batch_id][tensor_id] for batch_id in range(batch_num)], dim=dim_dict[key]) |
| 156 | out_dict[key] = feature_list |
| 157 | |
| 158 | return out_dict |
| 159 | |
| 160 | |
| 161 | def reshape( |
no test coverage detected