File size: 2,008 Bytes
cb9e677 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import contextlib
import dataclasses
import datetime
import logging
import time
from typing import Optional, Protocol
import torch
logger = logging.getLogger("utils")
@dataclasses.dataclass
class TrainState:
max_steps: int
step: int = 0
elapsed_time: float = 0.0
n_seen_tokens: int = 0
this_step_time: float = 0.0
begin_step_time: float = 0.0
this_eval_perplexity: Optional[float] = None
this_eval_loss: Optional[float] = None
def start_step(self):
self.step += 1
self.begin_step_time = time.time()
def end_step(self, n_batch_tokens: int):
self.this_step_time = time.time() - self.begin_step_time
self.this_step_tokens = n_batch_tokens
self.elapsed_time += self.this_step_time
self.n_seen_tokens += self.this_step_tokens
self.begin_step_time = time.time()
@property
def wps(self):
return self.this_step_tokens / self.this_step_time
@property
def avg_wps(self):
return self.n_seen_tokens / self.elapsed_time
@property
def eta(self):
steps_left = self.max_steps - self.step
avg_time_per_step = self.elapsed_time / self.step
return steps_left * avg_time_per_step
def set_random_seed(seed: int) -> None:
"""Set random seed for reproducibility."""
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
class Closable(Protocol):
def close(self):
pass
@contextlib.contextmanager
def logged_closing(thing: Closable, name: str):
"""
Logging the closing to be sure something is not hanging at exit time
"""
try:
setattr(thing, "wrapped_by_closing", True)
yield
finally:
logger.info(f"Closing: {name}")
try:
thing.close()
except Exception:
logger.error(f"Error while closing {name}!")
raise
logger.info(f"Closed: {name}")
def now_as_str() -> str:
return datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|