CarlosMalaga's picture
Upload 201 files
2f044c1 verified
raw
history blame
691 Bytes
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}")