File size: 1,374 Bytes
5fe8ec0 | 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 | import logging
logger = logging.getLogger(__name__)
import torch
import os
def print_rank(message):
"""If distributed is initialized, print the rank."""
if torch.distributed.is_initialized():
logger.info(f'rank{torch.distributed.get_rank()}: ' + message)
else:
logger.info(message)
def print_master(message):
"""If distributed is initialized print only on rank 0."""
if torch.distributed.is_initialized():
if torch.distributed.get_rank() == 0:
logger.info(message)
else:
logger.info(message)
def find_latest_checkpoint(output_dir):
""" Scan the output directory and return the latest checkpoint path """
if not os.path.exists(output_dir):
return None
checkpoints = [
os.path.join(output_dir, d) for d in os.listdir(output_dir)
if d.startswith("checkpoint-") and os.path.isdir(os.path.join(output_dir, d))
]
if not checkpoints:
return None
# Sort by checkpoint number and return the latest one
latest_checkpoint = max(checkpoints, key=lambda x: int(x.split("-")[-1]))
return latest_checkpoint
def batch_to_device(batch, device):
_batch = {}
for key, value in batch.items():
if isinstance(value, torch.Tensor):
_batch[key] = value.to(device)
else:
_batch[key] = value
return _batch
|