fix in dataloader. To long sequences are dropped lukas-blecher#8

Author: lukas-blecher

dataset/dataset.py

@@ -26,12 +26,12 @@
             [alb.ShiftScaleRotate(shift_limit=0, scale_limit=(-.15, 0), rotate_limit=1, border_mode=0, interpolation=3,
                                   value=[255, 255, 255], p=1),
              alb.GridDistortion(distort_limit=0.1, border_mode=0, interpolation=3, value=[255, 255, 255], p=.5)], p=.15),
-        alb.InvertImg(p=.15),
+        #alb.InvertImg(p=.15),
         alb.RGBShift(r_shift_limit=15, g_shift_limit=15,
                      b_shift_limit=15, p=0.3),
         alb.GaussNoise(10, p=.2),
         alb.RandomBrightnessContrast(.05, (-.2, 0), True, p=0.2),
-        alb.JpegCompression(95, p=.5),
+        alb.JpegCompression(95, p=.3),
         alb.ToGray(always_apply=True),
         alb.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
         # alb.Sharpen()
@@ -150,6 +150,13 @@ def prepare_data(self, batch):
         """
 
         eqs, ims = batch.T
+        tok = self.tokenizer(list(eqs), return_token_type_ids=False)
+        # pad with bos and eos token
+        for k, p in zip(tok, [[self.bos_token_id, self.eos_token_id], [1, 1]]):
+            tok[k] = pad_sequence([torch.LongTensor([p[0]]+x+[p[1]]) for x in tok[k]], batch_first=True, padding_value=self.pad_token_id)
+        # check if sequence length is too long
+        if self.max_seq_len < tok['attention_mask'].shape[1]:
+            return next(self)
         images = []
         for path in list(ims):
             im = cv2.imread(path)
@@ -162,13 +169,6 @@ def prepare_data(self, batch):
                 if np.random.random() < .04:
                     im[im != 255] = 0
             images.append(self.transform(image=im)['image'][:1])
-        tok = self.tokenizer(list(eqs), return_token_type_ids=False)
-        # pad with bos and eos token
-        for k, p in zip(tok, [[self.bos_token_id, self.eos_token_id], [1, 1]]):
-            tok[k] = pad_sequence([torch.LongTensor([p[0]]+x+[p[1]]) for x in tok[k]], batch_first=True, padding_value=self.pad_token_id)
-        # check if sequence length is too long
-        if self.max_seq_len < len(tok[0]):
-            return next(self)
         try:
             images = torch.cat(images).float().unsqueeze(1)
         except RuntimeError:

train.py

@@ -51,7 +51,6 @@ def save_models(e):
                     torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
                     opt.step()
                     scheduler.step()
-
                     dset.set_description('Loss: %.4f' % loss.item())
                     if args.wandb:
                         wandb.log({'train/loss': loss.item()})
@@ -79,8 +78,7 @@ def save_models(e):
     parsed_args = parser.parse_args()
     with parsed_args.config as f:
         params = yaml.load(f, Loader=yaml.FullLoader)
-    args = parse_args(Munch(params))
-
+    args = parse_args(Munch(params), **vars(parsed_args))
     logging.getLogger().setLevel(logging.DEBUG if parsed_args.debug else logging.WARNING)
     seed_everything(args.seed)
     if args.wandb: