数据增强

imitate_episodes.py

@@ -154,6 +154,7 @@ def main(args):
         text_feature_dim=text_feature_dim,
         text_tokenizer_name=text_tokenizer_name,
         text_max_length=text_max_length,
+        image_augment=args['image_aug'],
     )
 
     # save dataset stats
@@ -516,5 +517,7 @@ if __name__ == '__main__':
     parser.add_argument('--text_encoder_type', action='store', type=str, required=False)
     parser.add_argument('--freeze_text_encoder', action='store_true')
     parser.add_argument('--text_max_length', action='store', type=int, required=False)
+    parser.add_argument('--image_aug', action='store_true',
+                        help='Enable training-time image augmentation (color/highlight/noise/blur)')
     
     main(vars(parser.parse_args()))