x-transformer decoder args in config

Author: lukas-blecher

models.py

@@ -129,7 +129,7 @@ def get_model(args):
                 dim=args.dim,
                 depth=args.num_layers,
                 heads=args.heads,
-                cross_attend=True
+                **args.decoder_args
             )),
         pad_value=args.pad_token
     ).to(args.device)

settings/config.yaml

@@ -2,6 +2,8 @@ batchsize: 10
 bos_token: 1
 channels: 1
 debug: false
+decoder_args: 
+  cross_attend: true
 device: cuda
 dim: 256
 encoder_depth: 4

settings/default.yaml

@@ -37,6 +37,12 @@ backbone_layers:
   - 9
 encoder_depth: 4
 num_layers: 4
+decoder_args: 
+  cross_attend: true
+  ff_glu: true
+  attn_on_attn: false
+  use_scalenorm: true
+  rel_pos_bias: false
 heads: 8
 num_tokens: 8000
 max_seq_len: 1024

utils/utils.py

@@ -54,6 +54,8 @@ def parse_args(args, **kwargs):
     args.wandb = not kwargs.debug and not args.debug
     args.device = 'cuda' if torch.cuda.is_available() and not kwargs.no_cuda else 'cpu'
     args.max_dimensions = [args.max_width, args.max_height]
+    if 'decoder_args' not in args or args.decoder_args is None:
+        args.decoder_args = {}
     if 'model_path' in args:
         args.out_path = os.path.join(args.model_path, args.name)
         os.makedirs(args.out_path, exist_ok=True)