validate with custom temperature

Author: lukas-blecher

eval.py

@@ -12,7 +12,7 @@
 from tqdm.auto import tqdm
 import wandb
 
-from models import get_model
+from models import get_model, Model
 from utils import *
 
 
@@ -29,7 +29,7 @@ def detokenize(tokens, tokenizer):
 
 
 @torch.no_grad()
-def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, args: Munch, num_batches: int = None, name: str = 'test'):
+def evaluate(model: Model, dataset: Im2LatexDataset, args: Munch, num_batches: int = None, name: str = 'test'):
     """evaluates the model. Returns bleu score on the dataset
 
     Args:
@@ -53,7 +53,7 @@ def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, args: Munch, num_
         encoded = model.encoder(im.to(device))
         #loss = decoder(tgt_seq, mask=tgt_mask, context=encoded)
         dec = model.decoder.generate(torch.LongTensor([args.bos_token]*len(encoded))[:, None].to(device), args.max_seq_len,
-                                     eos_token=args.pad_token, context=encoded)
+                                     eos_token=args.pad_token, context=encoded, temperature=(args.temperature if 'temperature' in args else 1))
         pred = detokenize(dec, dataset.tokenizer)
         truth = detokenize(seq['input_ids'], dataset.tokenizer)
         bleus.append(metrics.bleu_score(pred, [alternatives(x) for x in truth]))
@@ -84,13 +84,15 @@ def evaluate(model: torch.nn.Module, dataset: Im2LatexDataset, args: Munch, num_
     parser.add_argument('--no-cuda', action='store_true', help='Use CPU')
     parser.add_argument('-b', '--batchsize', type=int, default=10, help='Batch size')
     parser.add_argument('--debug', action='store_true', help='DEBUG')
+    parser.add_argument('-t', '--temperature', type=float, default=.333, help='sampling emperature')
 
     parsed_args = parser.parse_args()
     with parsed_args.config as f:
         params = yaml.load(f, Loader=yaml.FullLoader)
     args = parse_args(Munch(params))
     args.testbatchsize = parsed_args.batchsize
     args.wandb = False
+    args.temperature = parsed_args.temperature
     logging.getLogger().setLevel(logging.DEBUG if parsed_args.debug else logging.WARNING)
     seed_everything(args.seed if 'seed' in args else 42)
     model = get_model(args)

models.py

@@ -10,24 +10,6 @@
 from einops import rearrange, repeat
 
 
-class Model(nn.Module):
-    def __init__(self, encoder: Encoder, decoder: AutoregressiveWrapper, args, temp: float = .333):
-        super().__init__()
-        self.encoder = encoder
-        self.decoder = decoder
-        self.bos_token = args.bos_token
-        self.eos_token = args.eos_token
-        self.max_seq_len = args.max_seq_len
-        self.temperature = temp
-
-    @torch.no_grad()
-    def forward(self, x: torch.Tensor):
-        device = x.device
-        encoded = self.encoder(x.to(device))
-        dec = self.decoder.generate(torch.LongTensor([self.bos_token]*len(x))[:, None].to(device), self.max_seq_len,
-                                    eos_token=self.eos_token, context=encoded, temperature=self.temperature)
-        return dec
-
 
 class CustomARWrapper(AutoregressiveWrapper):
     def __init__(self, *args, **kwargs):
@@ -106,6 +88,25 @@ def forward_features(self, x):
         return x
 
 
+class Model(nn.Module):
+    def __init__(self, encoder: CustomVisionTransformer, decoder: CustomARWrapper, args, temp: float = .333):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.bos_token = args.bos_token
+        self.eos_token = args.eos_token
+        self.max_seq_len = args.max_seq_len
+        self.temperature = temp
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor):
+        device = x.device
+        encoded = self.encoder(x.to(device))
+        dec = self.decoder.generate(torch.LongTensor([self.bos_token]*len(x))[:, None].to(device), self.max_seq_len,
+                                    eos_token=self.eos_token, context=encoded, temperature=self.temperature)
+        return dec
+
+
 def get_model(args):
     backbone = ResNetV2(
         layers=args.backbone_layers, num_classes=0, global_pool='', in_chans=args.channels,

settings/default.yaml

@@ -11,8 +11,11 @@ name: "pix2tex"
 # Training parameters
 epochs: 10
 batchsize: 8
+
+# Testing parameters
 testbatchsize: 20
 valbatches: 100
+temperature: 0.2
 
 # Optimizer configurations
 optimizer: "Adam"