hongshaorou/YOPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Speed up dataloader (train speed x110%)

Browse Source
This commit is contained in:
TJU_Lu 2025-07-16 21:41:27 +08:00
parent b7bfea3a5c
commit 6f7060838c
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
YOPO/policy/yopo_trainer.py
View File

@ -58,8 +58,10 @@ class YopoTrainer:
print("Network Loaded! Loading Dataset...")
# dataset
self.train_dataloader = DataLoader(YOPODataset(mode='train'), batch_size=self.batch_size, shuffle=True, num_workers=0)
self.val_dataloader = DataLoader(YOPODataset(mode='valid'), batch_size=self.batch_size, shuffle=False, num_workers=0)
self.train_dataloader = DataLoader(YOPODataset(mode='train'), batch_size=self.batch_size, shuffle=True,
num_workers=1, pin_memory=True)
self.val_dataloader = DataLoader(YOPODataset(mode='valid'), batch_size=self.batch_size, shuffle=False,
num_workers=1, pin_memory=True)
print("Dataset Loaded!")
def train(self, epoch, save_interval=None):