File size: 691 Bytes
2f044c1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
from typing import Dict, List
from collections import defaultdict
from lightning.pytorch.callbacks import Callback
from relik.common.log import get_logger
import os
logger = get_logger()
class ShuffleTrainCallback(Callback):
def __init__(self, shuffle_every: int = 1, data_path: str = None):
self.shuffle_every = shuffle_every
self.data_path = data_path
def on_train_epoch_end(self, trainer, pl_module):
if (trainer.current_epoch + 1) % self.shuffle_every == 0:
logger.info("Shuffling train dataset")
os.system(f"shuf {self.data_path} > {self.data_path}.shuf")
os.system(f"mv {self.data_path}.shuf {self.data_path}") |