()
| 47 | pass |
| 48 | |
| 49 | def train(): |
| 50 | utils.global_seed(cfg['seed']) # set global seed for numpy, random, pytorch, dgl, cuda and cudnn |
| 51 | print('cfg',cfg) |
| 52 | device = utils.set_gpu_mode(cfg['use_gpu'], cfg['gpu_id']) |
| 53 | print('device', device) |
| 54 | local_rank = cfg['train']['local_rank'] |
| 55 | rank = local_rank # for distributed training |
| 56 | logger, log_file, exp_id = create_logger(cfg, local_rank) |
| 57 | warnings.filterwarnings("ignore") |
| 58 | |
| 59 | auto_resume = cfg['train']['auto_resume'] # for resumed training |
| 60 | |
| 61 | # close loop |
| 62 | model_dir = osp.join(cfg['output_dir'], cfg['dataset']['dataset_name'], "models", exp_id) |
| 63 | |
| 64 | tensorboard_dir = ( |
| 65 | os.path.join(cfg['output_dir'], cfg['name'], "tensorboard") |
| 66 | if cfg['train']['tensorboard']['enable'] |
| 67 | else None |
| 68 | ) |
| 69 | |
| 70 | set_distributed(cfg=cfg, local_rank=local_rank, model_dir=model_dir, tensorboard_dir=tensorboard_dir) |
| 71 | |
| 72 | # ----- BEGIN DATASET BUILDER ----- |
| 73 | datasets = get_dataset(cfg) |
| 74 | train_set = datasets['train_set'] |
| 75 | valid_set = datasets['valid_set'] |
| 76 | entity_type = train_set.entity_type |
| 77 | |
| 78 | # ----- END DATASET BUILDER ----- |
| 79 | |
| 80 | # ----- BEGIN LOSS BUILDER ----- |
| 81 | if cfg['setting']['type'] not in ['LT Regression', 'LT Generation']: |
| 82 | num_class_list = get_category_list(train_set) |
| 83 | num_classes = len(num_class_list) |
| 84 | para_dict = { |
| 85 | "num_classes": num_classes, |
| 86 | "num_class_list": num_class_list, |
| 87 | "cfg": cfg, |
| 88 | "device": device, |
| 89 | } |
| 90 | cfg['setting']['num_class'] = num_classes # update the real number of classes based on the datasets |
| 91 | |
| 92 | criterion = eval(cfg['loss']['type'])(para_dict) |
| 93 | elif cfg['setting']['type'] == 'LT Regression': |
| 94 | para_dict = { |
| 95 | "cfg": cfg, |
| 96 | "device": device, |
| 97 | } |
| 98 | criterion = eval(cfg['loss']['type'])(para_dict) |
| 99 | num_class_list = None |
| 100 | else: |
| 101 | raise NotImplementedError |
| 102 | epoch_number = cfg['train']['max_epoch'] |
| 103 | |
| 104 | # ----- END LOSS BUILDER ----- |
| 105 | |
| 106 | # ----- BEGIN MODEL BUILDER ----- |
no test coverage detected