forked from mrq/ai-voice-cloning
fixes #134
This commit is contained in:
parent
4b952ea52a
commit
84b7383428
65
src/train.py
65
src/train.py
|
@ -6,35 +6,10 @@ import datetime
|
||||||
|
|
||||||
from torch.distributed.run import main as torchrun
|
from torch.distributed.run import main as torchrun
|
||||||
|
|
||||||
# I don't want this invoked from an import
|
|
||||||
if __name__ != "__main__":
|
|
||||||
raise Exception("Do not invoke this from an import")
|
|
||||||
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument('--yaml', type=str, help='Path to training configuration file.', default='./training/voice/train.yml', nargs='+') # ugh
|
|
||||||
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='Job launcher')
|
|
||||||
args = parser.parse_args()
|
|
||||||
args.yaml = " ".join(args.yaml) # absolutely disgusting
|
|
||||||
config_path = args.yaml
|
|
||||||
|
|
||||||
with open(config_path, 'r') as file:
|
|
||||||
opt_config = yaml.safe_load(file)
|
|
||||||
|
|
||||||
# it'd be downright sugoi if I was able to install DLAS as a pip package
|
# it'd be downright sugoi if I was able to install DLAS as a pip package
|
||||||
sys.path.insert(0, './modules/dlas/codes/')
|
sys.path.insert(0, './modules/dlas/codes/')
|
||||||
sys.path.insert(0, './modules/dlas/')
|
sys.path.insert(0, './modules/dlas/')
|
||||||
|
|
||||||
# yucky override
|
|
||||||
if "bitsandbytes" in opt_config and not opt_config["bitsandbytes"]:
|
|
||||||
os.environ['BITSANDBYTES_OVERRIDE_LINEAR'] = '0'
|
|
||||||
os.environ['BITSANDBYTES_OVERRIDE_EMBEDDING'] = '0'
|
|
||||||
os.environ['BITSANDBYTES_OVERRIDE_ADAM'] = '0'
|
|
||||||
os.environ['BITSANDBYTES_OVERRIDE_ADAMW'] = '0'
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from codes import train as tr
|
|
||||||
from utils import util, options as option
|
|
||||||
|
|
||||||
# this is effectively just copy pasted and cleaned up from the __main__ section of training.py
|
# this is effectively just copy pasted and cleaned up from the __main__ section of training.py
|
||||||
def train(config_path, launcher='none'):
|
def train(config_path, launcher='none'):
|
||||||
opt = option.parse(config_path, is_train=True)
|
opt = option.parse(config_path, is_train=True)
|
||||||
|
@ -59,13 +34,35 @@ def train(config_path, launcher='none'):
|
||||||
trainer.init(config_path, opt, launcher, '')
|
trainer.init(config_path, opt, launcher, '')
|
||||||
trainer.do_training()
|
trainer.do_training()
|
||||||
|
|
||||||
try:
|
if __name__ == "__main__":
|
||||||
import torch_intermediary
|
parser = argparse.ArgumentParser()
|
||||||
if torch_intermediary.OVERRIDE_ADAM:
|
parser.add_argument('--yaml', type=str, help='Path to training configuration file.', default='./training/voice/train.yml', nargs='+') # ugh
|
||||||
print("Using BitsAndBytes optimizations")
|
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='Job launcher')
|
||||||
else:
|
args = parser.parse_args()
|
||||||
print("NOT using BitsAndBytes optimizations")
|
args.yaml = " ".join(args.yaml) # absolutely disgusting
|
||||||
except Exception as e:
|
config_path = args.yaml
|
||||||
pass
|
|
||||||
|
|
||||||
train(config_path, args.launcher)
|
with open(config_path, 'r') as file:
|
||||||
|
opt_config = yaml.safe_load(file)
|
||||||
|
|
||||||
|
# yucky override
|
||||||
|
if "bitsandbytes" in opt_config and not opt_config["bitsandbytes"]:
|
||||||
|
os.environ['BITSANDBYTES_OVERRIDE_LINEAR'] = '0'
|
||||||
|
os.environ['BITSANDBYTES_OVERRIDE_EMBEDDING'] = '0'
|
||||||
|
os.environ['BITSANDBYTES_OVERRIDE_ADAM'] = '0'
|
||||||
|
os.environ['BITSANDBYTES_OVERRIDE_ADAMW'] = '0'
|
||||||
|
|
||||||
|
try:
|
||||||
|
import torch_intermediary
|
||||||
|
if torch_intermediary.OVERRIDE_ADAM:
|
||||||
|
print("Using BitsAndBytes optimizations")
|
||||||
|
else:
|
||||||
|
print("NOT using BitsAndBytes optimizations")
|
||||||
|
except Exception as e:
|
||||||
|
pass
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from codes import train as tr
|
||||||
|
from utils import util, options as option
|
||||||
|
|
||||||
|
train(config_path, args.launcher)
|
Loading…
Reference in New Issue
Block a user