(tensor_set, chunk=None)
| 53 | |
| 54 | |
| 55 | def flatten_tensors(tensor_set, chunk=None): |
| 56 | tensors = [p for p in tensor_set] |
| 57 | weights = [p.data for p in tensors] |
| 58 | sizes = [p.numel() for p in tensors] |
| 59 | total_size = sum(sizes) |
| 60 | if chunk: |
| 61 | total_size = ((total_size+chunk-1)//chunk)*chunk |
| 62 | |
| 63 | flatten_weights_tensor = torch.zeros(total_size, dtype=weights[0].dtype).to(weights[0].device) |
| 64 | flatten_weights_storage = flatten_weights_tensor.storage() |
| 65 | |
| 66 | def set_storage(param, weight_storage, storage_offset): |
| 67 | with torch.no_grad(): |
| 68 | z = torch.zeros_like(param.data) |
| 69 | z.set_(weight_storage, storage_offset, param.shape) |
| 70 | param.data = z |
| 71 | |
| 72 | offset = 0 |
| 73 | for i in range(len(tensors)): |
| 74 | flatten_weights_tensor[offset: offset + weights[i].numel()] = weights[i].reshape(-1) |
| 75 | set_storage(tensors[i], flatten_weights_storage, offset) |
| 76 | offset += sizes[i] |
| 77 | |
| 78 | return flatten_weights_tensor |
nothing calls this directly
no test coverage detected