File size: 1,951 Bytes
3b99abb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import sys
sys.path.append('.')
import yaml
import argparse
from easydict import EasyDict
from scripts.utils.others import setup_seed
from scripts.utils.module_loader import *
def run(config):
# Initialize a model
model = load_model(config.model)
# for i, (name, param) in enumerate(model.named_parameters()):
# print(f"{i}: {name}", param.requires_grad, id(param))
# return
# Initialize a dataset
data_module = load_dataset(config.dataset)
# Initialize a trainer
trainer = load_trainer(config)
# Train and validate
trainer.fit(model=model, datamodule=data_module)
# Load best model and test performance
if model.save_path is not None:
if config.model.kwargs.get("use_lora", False):
# Load LoRA model
config.model.kwargs.lora_config_path = model.save_path
model = load_model(config.model)
else:
model.load_checkpoint(model.save_path, load_prev_scheduler=model.load_prev_scheduler)
trainer.test(model=model, datamodule=data_module)
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--config', help="running configurations", type=str, required=True)
return parser.parse_args()
def main(args):
with open(args.config, 'r', encoding='utf-8') as r:
config = EasyDict(yaml.safe_load(r))
if config.setting.seed:
setup_seed(config.setting.seed)
# set os environment variables
for k, v in config.setting.os_environ.items():
if v is not None and k not in os.environ:
os.environ[k] = str(v)
elif k in os.environ:
# override the os environment variables
config.setting.os_environ[k] = os.environ[k]
# Only the root node will print the log
if config.setting.os_environ.NODE_RANK != 0:
config.Trainer.logger = False
run(config)
if __name__ == '__main__':
main(get_args()) |