| 165 | |
| 166 | |
| 167 | class TensorboardLogger(object): |
| 168 | def __init__(self, log_dir): |
| 169 | self.writer = SummaryWriter(logdir=log_dir) |
| 170 | self.step = 0 |
| 171 | |
| 172 | def set_step(self, step=None): |
| 173 | if step is not None: |
| 174 | self.step = step |
| 175 | else: |
| 176 | self.step += 1 |
| 177 | |
| 178 | def update(self, head='scalar', step=None, **kwargs): |
| 179 | for k, v in kwargs.items(): |
| 180 | if v is None: |
| 181 | continue |
| 182 | if isinstance(v, torch.Tensor): |
| 183 | v = v.item() |
| 184 | assert isinstance(v, (float, int)) |
| 185 | self.writer.add_scalar(head + "/" + k, v, self.step if step is None else step) |
| 186 | |
| 187 | def flush(self): |
| 188 | self.writer.flush() |
| 189 | |
| 190 | def seed_worker(worker_id): |
| 191 | worker_seed = torch.initial_seed() % 2**32 |
nothing calls this directly
no outgoing calls
no test coverage detected