|
|
from typing import Dict, Callable, List |
|
|
import collections |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
def dict_apply( |
|
|
x: Dict[str, torch.Tensor], |
|
|
func: Callable[[torch.Tensor], torch.Tensor] |
|
|
) -> Dict[str, torch.Tensor]: |
|
|
result = dict() |
|
|
for key, value in x.items(): |
|
|
if isinstance(value, dict): |
|
|
result[key] = dict_apply(value, func) |
|
|
else: |
|
|
result[key] = func(value) |
|
|
return result |
|
|
|
|
|
def pad_remaining_dims(x, target): |
|
|
assert x.shape == target.shape[:len(x.shape)] |
|
|
return x.reshape(x.shape + (1,)*(len(target.shape) - len(x.shape))) |
|
|
|
|
|
def dict_apply_split( |
|
|
x: Dict[str, torch.Tensor], |
|
|
split_func: Callable[[torch.Tensor], Dict[str, torch.Tensor]] |
|
|
) -> Dict[str, torch.Tensor]: |
|
|
results = collections.defaultdict(dict) |
|
|
for key, value in x.items(): |
|
|
result = split_func(value) |
|
|
for k, v in result.items(): |
|
|
results[k][key] = v |
|
|
return results |
|
|
|
|
|
def dict_apply_reduce( |
|
|
x: List[Dict[str, torch.Tensor]], |
|
|
reduce_func: Callable[[List[torch.Tensor]], torch.Tensor] |
|
|
) -> Dict[str, torch.Tensor]: |
|
|
result = dict() |
|
|
for key in x[0].keys(): |
|
|
result[key] = reduce_func([x_[key] for x_ in x]) |
|
|
return result |
|
|
|
|
|
|
|
|
def replace_submodules( |
|
|
root_module: nn.Module, |
|
|
predicate: Callable[[nn.Module], bool], |
|
|
func: Callable[[nn.Module], nn.Module]) -> nn.Module: |
|
|
""" |
|
|
predicate: Return true if the module is to be replaced. |
|
|
func: Return new module to use. |
|
|
""" |
|
|
if predicate(root_module): |
|
|
return func(root_module) |
|
|
|
|
|
bn_list = [k.split('.') for k, m |
|
|
in root_module.named_modules(remove_duplicate=True) |
|
|
if predicate(m)] |
|
|
for *parent, k in bn_list: |
|
|
parent_module = root_module |
|
|
if len(parent) > 0: |
|
|
parent_module = root_module.get_submodule('.'.join(parent)) |
|
|
if isinstance(parent_module, nn.Sequential): |
|
|
src_module = parent_module[int(k)] |
|
|
else: |
|
|
src_module = getattr(parent_module, k) |
|
|
tgt_module = func(src_module) |
|
|
if isinstance(parent_module, nn.Sequential): |
|
|
parent_module[int(k)] = tgt_module |
|
|
else: |
|
|
setattr(parent_module, k, tgt_module) |
|
|
|
|
|
bn_list = [k.split('.') for k, m |
|
|
in root_module.named_modules(remove_duplicate=True) |
|
|
if predicate(m)] |
|
|
assert len(bn_list) == 0 |
|
|
return root_module |
|
|
|
|
|
def optimizer_to(optimizer, device): |
|
|
for state in optimizer.state.values(): |
|
|
for k, v in state.items(): |
|
|
if isinstance(v, torch.Tensor): |
|
|
state[k] = v.to(device=device) |
|
|
return optimizer |
|
|
|
