| 2 | |
| 3 | |
| 4 | class ForwardHookForDevice: |
| 5 | def __init__(self): |
| 6 | pass |
| 7 | |
| 8 | @staticmethod |
| 9 | def get_align_device_pre_forward_hook(device="cuda", with_kwargs=False): |
| 10 | """ |
| 11 | ensure same device for input and module |
| 12 | """ |
| 13 | |
| 14 | def hook(module: torch.nn.Module, args): |
| 15 | if device is not None: |
| 16 | align_device = device |
| 17 | elif len(list(module.parameters())) > 0: |
| 18 | align_device = next(module.parameters()).device |
| 19 | else: |
| 20 | align_device = "cuda" |
| 21 | module.to(align_device) |
| 22 | args = tuple(arg.to(align_device) if isinstance(arg, torch.Tensor) else arg for arg in args) |
| 23 | return args |
| 24 | |
| 25 | def hook_with_kwargs(module: torch.nn.Module, args, kwargs): |
| 26 | if device is not None: |
| 27 | align_device = device |
| 28 | elif len(list(module.parameters())) > 0: |
| 29 | align_device = next(module.parameters()).device |
| 30 | else: |
| 31 | align_device = "cuda" |
| 32 | module.to(align_device) |
| 33 | args = tuple(arg.to(align_device) if isinstance(arg, torch.Tensor) else arg for arg in args) |
| 34 | _kwargs = dict() |
| 35 | for k, v in kwargs.items(): |
| 36 | if isinstance(v, torch.Tensor): |
| 37 | _kwargs[k] = v.to(align_device) |
| 38 | else: |
| 39 | _kwargs[k] = v |
| 40 | kwargs = _kwargs |
| 41 | return args, kwargs |
| 42 | |
| 43 | if with_kwargs: |
| 44 | return hook_with_kwargs |
| 45 | else: |
| 46 | return hook |
| 47 | |
| 48 | @staticmethod |
| 49 | def get_forward_hook(pre: bool, device=None, with_kwargs=False): |
| 50 | """ |
| 51 | device is executing device |
| 52 | origin_device is the device where tensor is saved after forward |
| 53 | """ |
| 54 | origin_device = "cpu" |
| 55 | if device is not None: |
| 56 | device = device |
| 57 | else: |
| 58 | device = "cuda" |
| 59 | |
| 60 | def pre_hook(module: torch.nn.Module, args): |
| 61 | module.to(device) |
nothing calls this directly
no outgoing calls
no test coverage detected