compatible with torch.utils.data.DataLoader

pix2tex/dataset/dataset.py

@@ -243,7 +243,7 @@ def update(self, **kwargs):
 class Dataloader(DataLoader):
  def __init__(self, dataset: Im2LatexDataset, batch_size=1, shuffle=False, *args, **kwargs):
  self.dataset = dataset
- self.dataset.update(batchsize=batch_size, shuffle=shuffle)
+ self.dataset.update(batchsize=batch_size, shuffle=shuffle, *args, **kwargs)
  super().__init__(self.dataset, *args, shuffle=False, batch_size=None, **kwargs)
 
  def __iter__(self):

pix2tex/eval.py

@@ -1,4 +1,4 @@
-from pix2tex.dataset.dataset import Im2LatexDataset
+from pix2tex.dataset.dataset import Im2LatexDataset, Dataloader
 import argparse
 import logging
 import yaml
@@ -28,12 +28,12 @@ def detokenize(tokens, tokenizer):
 
 
 @torch.no_grad()
-def evaluate(model: Model, dataset: Im2LatexDataset, args: Munch, num_batches: int = None, name: str = 'test'):
+def evaluate(model: Model, dataset: Dataloader, args: Munch, num_batches: int = None, name: str = 'test'):
  """evaluates the model. Returns bleu score on the dataset
 
  Args:
  model (torch.nn.Module): the model
- dataset (Im2LatexDataset): test dataset
+ dataset (Dataloader): test dataset
  args (Munch): arguments
  num_batches (int): How many batches to evaluate on. Defaults to None (all batches).
  name (str, optional): name of the test e.g. val or test for wandb. Defaults to 'test'.

pix2tex/model/settings/config-vit.yaml

@@ -1,4 +1,5 @@
 gpu_devices: null #[0,1,2,3,4,5,6,7]
+num_workers: 0
 betas:
 - 0.9
 - 0.999

pix2tex/model/settings/config.yaml

@@ -1,4 +1,5 @@
 gpu_devices: null #[0,1,2,3,4,5,6,7]
+num_workers: 0
 backbone_layers:
 - 2
 - 3

pix2tex/model/settings/debug.yaml

@@ -65,3 +65,7 @@ pad: False
 pad_token: 0
 bos_token: 1
 eos_token: 2
+
+#devices(GPU&CPU)
+num_workers: 0
+gpu_devices: null #[0,1,2,3,4,5,6,7]

pix2tex/train.py

@@ -1,4 +1,4 @@
-from pix2tex.dataset.dataset import Im2LatexDataset
+from pix2tex.dataset.dataset import Im2LatexDataset, Dataloader
 import os
 import argparse
 import logging
@@ -16,12 +16,12 @@
 
 
 def train(args):
- dataloader = Im2LatexDataset().load(args.data)
- dataloader.update(**args, test=False)
- valdataloader = Im2LatexDataset().load(args.valdata)
+ train_dataset = Im2LatexDataset().load(args.data)
+ train_dataloader = Dataloader(train_dataset, **args, test=False)
+ val_dataset = Im2LatexDataset().load(args.valdata)
  valargs = args.copy()
  valargs.update(batchsize=args.testbatchsize, keep_smaller_batches=True, test=True)
- valdataloader.update(**valargs)
+ val_dataloader = Dataloader(val_dataset, **valargs)
  device = args.device
  model = get_model(args)
  if torch.cuda.is_available() and not args.no_cuda:
@@ -47,7 +47,7 @@ def save_models(e, step=0):
  try:
  for e in range(args.epoch, args.epochs):
  args.epoch = e
- dset = tqdm(iter(dataloader))
+ dset = tqdm(iter(train_dataloader))
  for i, (seq, im) in enumerate(dset):
  if seq is not None and im is not None:
  opt.zero_grad()
@@ -63,20 +63,20 @@ def save_models(e, step=0):
  dset.set_description('Loss: %.4f' % total_loss)
  if args.wandb:
  wandb.log({'train/loss': total_loss})
- if (i+1+len(dataloader)*e) % args.sample_freq == 0:
- bleu_score, edit_distance, token_accuracy = evaluate(model, valdataloader, args, num_batches=int(args.valbatches*e/args.epochs), name='val')
+ if (i+1+len(train_dataloader)*e) % args.sample_freq == 0:
+ bleu_score, edit_distance, token_accuracy = evaluate(model, val_dataloader, args, num_batches=int(args.valbatches*e/args.epochs), name='val')
  if bleu_score > max_bleu and token_accuracy > max_token_acc:
  max_bleu, max_token_acc = bleu_score, token_accuracy
  save_models(e, step=i)
  if (e+1) % args.save_freq == 0:
- save_models(e, step=len(dataloader))
+ save_models(e, step=len(train_dataloader))
  if args.wandb:
  wandb.log({'train/epoch': e+1})
  except KeyboardInterrupt:
  if e >= 2:
  save_models(e, step=i)
  raise KeyboardInterrupt
- save_models(e, step=len(dataloader))
+ save_models(e, step=len(train_dataloader))
 
 
 if __name__ == '__main__':

pix2tex/utils/utils.py

@@ -55,6 +55,7 @@ def parse_args(args, **kwargs) -> Munch:
  args.update(kwargs)
  args.wandb = not kwargs.debug and not args.debug
  args.device = get_device(args, kwargs.no_cuda)
+ args.num_workers = args.get('num_workers', 0)
  args.max_dimensions = [args.max_width, args.max_height]
  args.min_dimensions = [args.get('min_width', 32), args.get('min_height', 32)]
  if 'decoder_args' not in args or args.decoder_args is None: