stop eval after n batches

+weird import error

Author: lukas-blecher

dataset/dataset.py

@@ -1,3 +1,5 @@
+import albumentations as alb
+from albumentations.pytorch import ToTensorV2
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -15,8 +17,6 @@
 import cv2
 from transformers import PreTrainedTokenizerFast
 from tqdm.auto import tqdm
-import albumentations as alb
-from albumentations.pytorch import ToTensorV2
 
 
 train_transform = alb.Compose(

eval.py

@@ -27,13 +27,14 @@ def detokenize(tokens, tokenizer):
 
 
 @torch.no_grad()
-def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, args: Munch, name: str = 'test'):
+def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, 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
         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'.
 
     Returns:
@@ -53,6 +54,8 @@ def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, args: Munch, name
         truth = detokenize(seq['input_ids'], dataset.tokenizer)
         bleus.append(metrics.bleu_score(pred, [alternatives(x) for x in truth]))
         pbar.set_description('BLEU: %.2f' % (np.mean(bleus)))
+        if num_batches is not None and i >= num_batches:
+            break
     bleu_score = np.mean(bleus)
     # samples
     pred = token2str(dec, dataset.tokenizer)

settings/default.yaml

@@ -12,6 +12,7 @@ name: "pix2tex"
 epochs: 10
 batchsize: 8
 testbatchsize: 20
+valbatches: 100
 
 # Optimizer configurations
 optimizer: "Adam"

train.py

@@ -51,7 +51,7 @@ def train(args):
             if args.wandb:
                 wandb.log({'train/loss': loss.item()})
             if (i+1) % args.sample_freq == 0:
-                evaluate(model, valdataloader, args, name='val')
+                evaluate(model, valdataloader, args, num_batches=args.valbatches, name='val')
         if (e+1) % args.save_freq == 0:
             torch.save(model.state_dict(), os.path.join(args.out_path, '%s_e%02d.pth' % (args.name, e+1)))
             yaml.dump(dict(args), open(os.path.join(args.out_path, 'config.yaml'), 'w+'))