github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

chen pcie profiler `to_tensorboard`

pull/717/head
Jie Zhu 2022-04-11 14:28:02 +08:00
parent 1e62130ce2
commit cd9492854b
No known key found for this signature in database
GPG Key ID: 58B7CA9DB6DC5BD9
1 changed files with 21 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
colossalai/utils/profiler/pcie_profiler.py
View File

@ -2,7 +2,7 @@ from pathlib import Path
from torch.autograd.profiler import profile from torch.autograd.profiler import profile
from .prof_utils import BaseProfiler, _format_time, _format_memory, _format_bandwidth from .prof_utils import BaseProfiler, _format_time, _format_memory, _format_bandwidth
from typing import List from typing import List
import json
def _get_size(dtype: str): def _get_size(dtype: str):
if dtype == "fp16": if dtype == "fp16":
@ -104,8 +104,26 @@ class PcieProfiler(BaseProfiler):
self.profiler = None self.profiler = None
def to_tensorboard(self, writer): def to_tensorboard(self, json_dir: Path):
writer.add_text(tag="Data Transmission", text_string=self.result_str("\n\n")) data = {
"h2d_time": self.h2d_time,
"h2d_count": self.h2d_count,
"d2h_time": self.d2h_time,
"d2h_count": self.d2h_count,
}
events_list = []
for location, event in self.ops_record.items():
events_list.append({
"location": location.splitlines(),
"cuda_time": event.cuda_time,
"pcie_vol": event.pcie_vol,
"count": event.count
})
data["events"] = events_list
with open(json_dir.joinpath("pcie.json"), "w") as f:
json.dump(data, f)
def to_file(self, filename: Path): def to_file(self, filename: Path):
with open(filename, "w") as f: with open(filename, "w") as f: