(module, *args, **kwargs)
| 96 | |
| 97 | @staticmethod |
| 98 | def _hook_rss_memory_post_forward(module, *args, **kwargs): |
| 99 | try: |
| 100 | import psutil |
| 101 | except (ImportError): |
| 102 | raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.") |
| 103 | |
| 104 | process = psutil.Process(os.getpid()) |
| 105 | mem = process.memory_info() |
| 106 | module.mem_rss_post_forward = mem.rss |
| 107 | mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward |
| 108 | module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0) |
| 109 | return None |
| 110 | |
| 111 | def add_memory_hooks(self): |
| 112 | """ Add a memory hook before and after each sub-module forward pass to record increase in memory consumption. |
nothing calls this directly
no outgoing calls
no test coverage detected