(param_set, chunk=None)
| 10 | |
| 11 | |
| 12 | def flatten_params(param_set, chunk=None): |
| 13 | params = [p for p in param_set] |
| 14 | weights = [p.data for p in params] |
| 15 | grads = [p.grad.data if p.grad is not None else torch.zeros_like(p.data) for p in params] |
| 16 | sizes = [p.numel() for p in params] |
| 17 | total_size = sum(sizes) |
| 18 | if chunk: |
| 19 | total_size = ((total_size+chunk-1)//chunk)*chunk |
| 20 | |
| 21 | flatten_weights_tensor = torch.zeros(total_size, dtype=weights[0].dtype).to(weights[0].device) |
| 22 | flatten_grads_tensor = torch.zeros(total_size, dtype=weights[0].dtype).to(weights[0].device) |
| 23 | flatten_weights_storage = flatten_weights_tensor.storage() |
| 24 | flatten_grads_storage = flatten_grads_tensor.storage() |
| 25 | |
| 26 | def set_storage(param, weight_storage, grad_storage, storage_offset): |
| 27 | with torch.no_grad(): |
| 28 | z = torch.zeros_like(param.data) |
| 29 | z.set_(weight_storage, storage_offset, param.shape) |
| 30 | param.data = z |
| 31 | |
| 32 | t = torch.zeros_like(param.data) |
| 33 | t.set_(grad_storage, storage_offset, param.shape) |
| 34 | param.grad = t |
| 35 | |
| 36 | offset = 0 |
| 37 | for i in range(len(params)): |
| 38 | flatten_weights_tensor[offset: offset + weights[i].numel()] = weights[i].reshape(-1) |
| 39 | flatten_grads_tensor[offset: offset + grads[i].numel()] = grads[i].reshape(-1) |
| 40 | set_storage(params[i], flatten_weights_storage, flatten_grads_storage, offset) |
| 41 | offset += sizes[i] |
| 42 | |
| 43 | weight_tensors = [p.data for p in params] |
| 44 | grad_tensors = [p.grad.data for p in params] |
| 45 | |
| 46 | _assert_contiguous(weight_tensors) |
| 47 | _assert_contiguous(grad_tensors) |
| 48 | |
| 49 | with torch.no_grad(): |
| 50 | flatten_para = torch.nn.Parameter(flatten_weights_tensor, requires_grad=False) |
| 51 | flatten_para.grad = flatten_grads_tensor |
| 52 | return flatten_para |
| 53 | |
| 54 | |
| 55 | def flatten_tensors(tensor_set, chunk=None): |
no test coverage detected