better memory check

also more consistent validation frequency

pix2tex/models.py

@@ -146,10 +146,14 @@ def embed_layer(**x):
     model = Model(encoder, decoder, args)
     if training:
         # check if largest batch can be handled by system
-        batchsize = args.batchsize if args.get('micro_batchsize', -1) == -1 else args.micro_batchsize
-        im = torch.empty(batchsize, args.channels, args.max_height, args.min_height, device=args.device).float()
-        seq = torch.randint(0, args.num_tokens, (batchsize, args.max_seq_len), device=args.device).long()
-        decoder(seq, context=encoder(im)).sum().backward()
+        try:
+            batchsize = args.batchsize if args.get('micro_batchsize', -1) == -1 else args.micro_batchsize
+            for _ in range(5):
+                im = torch.empty(batchsize, args.channels, args.max_height, args.min_height, device=args.device).float()
+                seq = torch.randint(0, args.num_tokens, (batchsize, args.max_seq_len), device=args.device).long()
+                decoder(seq, context=encoder(im)).sum().backward()
+        except RuntimeError:
+            raise RuntimeError("The system cannot handle a batch size of %i for the maximum image size (%i, %i). Try to use a smaller micro batchsize."%(batchsize, args.max_height, args.max_width))
         model.zero_grad()
         torch.cuda.empty_cache()
         del im, seq

pix2tex/train.py

@@ -60,7 +60,7 @@ def save_models(e):
                     dset.set_description('Loss: %.4f' % total_loss)
                     if args.wandb:
                         wandb.log({'train/loss': total_loss})
-                if (i+1) % args.sample_freq == 0:
+                if (i+1+len(dataloader)*e) % args.sample_freq == 0:
                     evaluate(model, valdataloader, args, num_batches=int(args.valbatches*e/args.epochs), name='val')
             if (e+1) % args.save_freq == 0:
                 save_models(e)