MCPcopy Create free account
hub / github.com/DrugLT/ImDrug / train

Function train

script/train.py:49–281  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

47 pass
48
49def 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 -----

Callers 1

train.pyFile · 0.85

Calls 12

create_loggerFunction · 0.90
set_distributedFunction · 0.90
get_datasetFunction · 0.90
get_category_listFunction · 0.90
get_modelFunction · 0.90
CombinerClass · 0.90
get_optimizerFunction · 0.90
get_schedulerFunction · 0.90
DataLoaderClass · 0.90
set_resumeFunction · 0.90
train_modelFunction · 0.90
valid_modelFunction · 0.90

Tested by

no test coverage detected