Added AdamW 8 bit optimizer which can save some GPU memory during tra…

…ining (use 'adamw8bit' for training.optimizer in config). Also added possibility to provide optimizer parameters in config file using keyword 'optimizer'.

requirements.txt

@@ -23,3 +23,4 @@ torch_audiomentations
 asteroid==0.7.0
 auraloss
 torchseg
+bitsandbytes

train.py

@@ -395,15 +395,23 @@ def train_model(args):
     if 0:
         valid_multi_gpu(model, args, config, verbose=True)
 
+    optim_params = dict()
+    if 'optimizer' in config:
+        optim_params = dict(config['optimizer'])
+        print('Optimizer params from config:\n{}'.format(optim_params))
+
     if config.training.optimizer == 'adam':
-        optimizer = Adam(model.parameters(), lr=config.training.lr)
+        optimizer = Adam(model.parameters(), lr=config.training.lr, **optim_params)
     elif config.training.optimizer == 'adamw':
-        optimizer = AdamW(model.parameters(), lr=config.training.lr)
+        optimizer = AdamW(model.parameters(), lr=config.training.lr, **optim_params)
     elif config.training.optimizer == 'radam':
-        optimizer = RAdam(model.parameters(), lr=config.training.lr)
+        optimizer = RAdam(model.parameters(), lr=config.training.lr, **optim_params)
+    elif config.training.optimizer == 'adamw8bit':
+        import bitsandbytes as bnb
+        optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=config.training.lr, **optim_params)
     elif config.training.optimizer == 'sgd':
         print('Use SGD optimizer')
-        optimizer = SGD(model.parameters(), lr=config.training.lr, momentum=0.999)
+        optimizer = SGD(model.parameters(), lr=config.training.lr, **optim_params)
     else:
         print('Unknown optimizer: {}'.format(config.training.optimizer))
         exit()