def run_recbole(model=None, dataset=None, config_file_list=None, config_dict=None, saved=True):
r""" A fast running api, which includes the complete process of
training and testing a model on a specified dataset
Args:
model (str, optional): Model name. Defaults to ``None``.
dataset (str, optional): Dataset name. Defaults to ``None``.
config_file_list (list, optional): Config files used to modify experiment parameters. Defaults to ``None``.
config_dict (dict, optional): Parameters dictionary used to modify experiment parameters. Defaults to ``None``.
saved (bool, optional): Whether to save the model. Defaults to ``True``.
"""
# configurations initialization
config = Config(model=model, dataset=dataset, config_file_list=config_file_list, config_dict=config_dict)
init_seed(config['seed'], config['reproducibility'])
# logger initialization
init_logger(config)
logger = getLogger()
logger.info(config)
# dataset filtering
dataset = create_dataset(config)
if config['save_dataset']:
dataset.save()
logger.info(dataset)
# dataset splitting
train_data, valid_data, test_data = data_preparation(config, dataset)
if config['save_dataloaders']:
save_split_dataloaders(config, dataloaders=(train_data, valid_data, test_data))
# model loading and initialization
model = get_model(config['model'])(config, train_data.dataset).to(config['device'])
logger.info(model)
# trainer loading and initialization
trainer = get_trainer(config['MODEL_TYPE'], config['model'])(config, model)
# model training
best_valid_score, best_valid_result = trainer.fit(
train_data, valid_data, saved=saved, show_progress=config['show_progress']
)
# model evaluation
test_result = trainer.evaluate(test_data, load_best_model=saved, show_progress=config['show_progress'])
logger.info(set_color('best valid ', 'yellow') + f': {best_valid_result}')
logger.info(set_color('test result', 'yellow') + f': {test_result}')
return {
'best_valid_score': best_valid_score,
'valid_score_bigger': config['valid_metric_bigger'],
'best_valid_result': best_valid_result,
'test_result': test_result
}
def fit(self,
train_data,
valid_data=None,
verbose=True,
saved=True,
show_progress=False):
for epoch_idx in range(self.epochs):
self._train_at_once(train_data, valid_data)
if (epoch_idx + 1) % self.eval_step == 0:
# evaluate
valid_start_time = time()
valid_score, valid_result = self._valid_epoch(valid_data)
self.best_valid_score, self.cur_step, stop_flag, update_flag = early_stopping(
valid_score,
self.best_valid_score,
self.cur_step,
max_step=self.stopping_step,
bigger=self.valid_metric_bigger)
valid_end_time = time()
valid_score_output = (set_color("epoch %d evaluating", 'green') + " [" + set_color("time", 'blue')
+ ": %.2fs, " + set_color("valid_score", 'blue') + ": %f]") % \
(epoch_idx, valid_end_time - valid_start_time, valid_score)
valid_result_output = set_color(
'valid result', 'blue') + ': \n' + dict2str(valid_result)
if verbose:
self.logger.info(valid_score_output)
self.logger.info(valid_result_output)
self.tensorboard.add_scalar('Vaild_score', valid_score,
epoch_idx)
if update_flag:
if saved:
self.model.save_model(self.temp_best_file)
self._save_checkpoint(epoch_idx)
update_output = set_color(
'Saving current best',
'blue') + ': %s' % self.saved_model_file
if verbose:
self.logger.info(update_output)
self.best_valid_result = valid_result
if stop_flag:
stop_output = 'Finished training, best eval result in epoch %d' % \
(epoch_idx - self.cur_step * self.eval_step)
if self.temp_file:
os.remove(self.temp_file)
if verbose:
self.logger.info(stop_output)
break
return self.best_valid_score, self.best_valid_result
def _train_epoch(self,
train_data,
epoch_idx,
loss_func=None,
show_progress=False):
r"""Train the model in an epoch
Args:
train_data (DataLoader): The train data.
epoch_idx (int): The current epoch id.
loss_func (function): The loss function of :attr:`model`. If it is ``None``, the loss function will be
:attr:`self.model.calculate_loss`. Defaults to ``None``.
show_progress (bool): Show the progress of training epoch. Defaults to ``False``.
Returns:
float/tuple: The sum of loss returned by all batches in this epoch. If the loss in each batch contains
multiple parts and the model return these multiple parts loss instead of the sum of loss, it will return a
tuple which includes the sum of loss in each part.
"""
self.model.train()
loss_func = loss_func or self.model.calculate_loss
total_loss = None
iter_data = (tqdm(
train_data,
total=len(train_data),
ncols=100,
desc=set_color(f"Train {epoch_idx:>5}", 'pink'),
) if show_progress else train_data)
for batch_idx, interaction in enumerate(iter_data):
interaction = interaction.to(self.device)
self.optimizer.zero_grad()
losses = loss_func(interaction)
if isinstance(losses, tuple):
loss = sum(losses)
loss_tuple = tuple(per_loss.item() for per_loss in losses)
total_loss = loss_tuple if total_loss is None else tuple(
map(sum, zip(total_loss, loss_tuple)))
else:
loss = losses
total_loss = losses.item(
) if total_loss is None else total_loss + losses.item()
self._check_nan(loss)
loss.backward()
if self.clip_grad_norm:
clip_grad_norm_(self.model.parameters(), **self.clip_grad_norm)
self.optimizer.step()
if self.gpu_available and show_progress:
iter_data.set_postfix_str(
set_color('GPU RAM: ' + get_gpu_usage(self.device),
'yellow'))
return total_loss
def _generate_train_loss_output(self, epoch_idx, s_time, e_time, losses):
des = self.config['loss_decimal_place'] or 4
train_loss_output = (set_color('epoch %d training', 'green') + ' [' +
set_color('time', 'blue') +
': %.2fs, ') % (epoch_idx, e_time - s_time)
if isinstance(losses, tuple):
des = (set_color('train_loss%d', 'blue') + ': %.' + str(des) + 'f')
train_loss_output += ', '.join(des % (idx + 1, loss)
for idx, loss in enumerate(losses))
else:
des = '%.' + str(des) + 'f'
train_loss_output += set_color('train loss',
'blue') + ': ' + des % losses
return train_loss_output + ']'
def data_preparation(config, dataset):
"""Split the dataset by :attr:`config['eval_args']` and create training, validation and test dataloader.
Note:
If we can load split dataloaders by :meth:`load_split_dataloaders`, we will not create new split dataloaders.
Args:
config (Config): An instance object of Config, used to record parameter information.
dataset (Dataset): An instance object of Dataset, which contains all interaction records.
Returns:
tuple:
- train_data (AbstractDataLoader): The dataloader for training.
- valid_data (AbstractDataLoader): The dataloader for validation.
- test_data (AbstractDataLoader): The dataloader for testing.
"""
dataloaders = load_split_dataloaders(config)
if dataloaders is not None:
train_data, valid_data, test_data = dataloaders
else:
model_type = config['MODEL_TYPE']
built_datasets = dataset.build()
train_dataset, valid_dataset, test_dataset = built_datasets
train_sampler, valid_sampler, test_sampler = create_samplers(config, dataset, built_datasets)
if model_type != ModelType.KNOWLEDGE:
train_data = get_dataloader(config, 'train')(config, train_dataset, train_sampler, shuffle=True)
else:
kg_sampler = KGSampler(dataset, config['train_neg_sample_args']['distribution'])
train_data = get_dataloader(config, 'train')(config, train_dataset, train_sampler, kg_sampler, shuffle=True)
valid_data = get_dataloader(config, 'evaluation')(config, valid_dataset, valid_sampler, shuffle=False)
test_data = get_dataloader(config, 'evaluation')(config, test_dataset, test_sampler, shuffle=False)
if config['save_dataloaders']:
save_split_dataloaders(config, dataloaders=(train_data, valid_data, test_data))
logger = getLogger()
logger.info(
set_color('[Training]: ', 'pink') + set_color('train_batch_size', 'cyan') + ' = ' +
set_color(f'[{config["train_batch_size"]}]', 'yellow') + set_color(' negative sampling', 'cyan') + ': ' +
set_color(f'[{config["neg_sampling"]}]', 'yellow')
)
logger.info(
set_color('[Evaluation]: ', 'pink') + set_color('eval_batch_size', 'cyan') + ' = ' +
set_color(f'[{config["eval_batch_size"]}]', 'yellow') + set_color(' eval_args', 'cyan') + ': ' +
set_color(f'[{config["eval_args"]}]', 'yellow')
)
return train_data, valid_data, test_data
def load_split_dataloaders(config):
"""Load split dataloaders if saved dataloaders exist and
their :attr:`config` of dataset are the same as current :attr:`config` of dataset.
Args:
config (Config): An instance object of Config, used to record parameter information.
Returns:
dataloaders (tuple of AbstractDataLoader or None): The split dataloaders.
"""
default_file = os.path.join(config['checkpoint_dir'], f'{config["dataset"]}-for-{config["model"]}-dataloader.pth')
dataloaders_save_path = config['dataloaders_save_path'] or default_file
if not os.path.exists(dataloaders_save_path):
return None
with open(dataloaders_save_path, 'rb') as f:
train_data, valid_data, test_data = pickle.load(f)
for arg in dataset_arguments + ['seed', 'repeatable', 'eval_args']:
if config[arg] != train_data.config[arg]:
return None
train_data.update_config(config)
valid_data.update_config(config)
test_data.update_config(config)
logger = getLogger()
logger.info(set_color('Load split dataloaders from', 'pink') + f': [{dataloaders_save_path}]')
return train_data, valid_data, test_data
def pretrain(self, train_data, verbose=True, show_progress=False):
for epoch_idx in range(self.start_epoch, self.pretrain_epochs):
# train
training_start_time = time()
train_loss = self._train_epoch(train_data,
epoch_idx,
show_progress=show_progress)
self.train_loss_dict[epoch_idx] = sum(train_loss) if isinstance(
train_loss, tuple) else train_loss
training_end_time = time()
train_loss_output = \
self._generate_train_loss_output(epoch_idx, training_start_time, training_end_time, train_loss)
if verbose:
self.logger.info(train_loss_output)
self._add_train_loss_to_tensorboard(epoch_idx, train_loss)
if (epoch_idx + 1) % self.save_step == 0:
saved_model_file = os.path.join(
self.checkpoint_dir,
'{}-{}-{}.pth'.format(self.config['model'],
self.config['dataset'],
str(epoch_idx + 1)))
self.save_pretrained_model(epoch_idx, saved_model_file)
update_output = set_color('Saving current',
'blue') + ': %s' % saved_model_file
if verbose:
self.logger.info(update_output)
return self.best_valid_score, self.best_valid_result
def _get_field_from_config(self):
super()._get_field_from_config()
self.head_entity_field = self.config['HEAD_ENTITY_ID_FIELD']
self.tail_entity_field = self.config['TAIL_ENTITY_ID_FIELD']
self.relation_field = self.config['RELATION_ID_FIELD']
self.entity_field = self.config['ENTITY_ID_FIELD']
self._check_field('head_entity_field', 'tail_entity_field',
'relation_field', 'entity_field')
self.set_field_property(self.entity_field, FeatureType.TOKEN,
FeatureSource.KG, 1)
self.logger.debug(
set_color('relation_field', 'blue') + f': {self.relation_field}')
self.logger.debug(
set_color('entity_field', 'blue') + f': {self.entity_field}')
def __str__(self):
"""
Model prints with number of trainable parameters
"""
model_parameters = filter(lambda p: p.requires_grad, self.parameters())
params = sum([np.prod(p.size()) for p in model_parameters])
return super().__str__() + set_color('\nTrainable parameters',
'blue') + f': {params}'
def _load_kg(self, token, dataset_path):
self.logger.debug(
set_color(f'Loading kg from [{dataset_path}].', 'green'))
kg_path = os.path.join(dataset_path, f'{token}.kg')
if not os.path.isfile(kg_path):
raise ValueError(f'[{token}.kg] not found in [{dataset_path}].')
df = self._load_feat(kg_path, FeatureSource.KG)
self._check_kg(df)
return df
def save_split_dataloaders(config, dataloaders):
"""Save split dataloaders.
Args:
config (Config): An instance object of Config, used to record parameter information.
dataloaders (tuple of AbstractDataLoader): The split dataloaders.
"""
save_path = config['checkpoint_dir']
saved_dataloaders_file = f'{config["dataset"]}-for-{config["model"]}-dataloader.pth'
file_path = os.path.join(save_path, saved_dataloaders_file)
logger = getLogger()
logger.info(
set_color('Saved split dataloaders', 'blue') + f': {file_path}')
with open(file_path, 'wb') as f:
pickle.dump(dataloaders, f)
def _load_link(self, token, dataset_path):
self.logger.debug(
set_color(f'Loading link from [{dataset_path}].', 'green'))
link_path = os.path.join(dataset_path, f'{token}.link')
if not os.path.isfile(link_path):
raise ValueError(f'[{token}.link] not found in [{dataset_path}].')
df = self._load_feat(link_path, 'link')
self._check_link(df)
item2entity, entity2item = {}, {}
for item_id, entity_id in zip(df[self.iid_field].values,
df[self.entity_field].values):
item2entity[item_id] = entity_id
entity2item[entity_id] = item_id
return item2entity, entity2item
def create_dataset(config):
"""Create dataset according to :attr:`config['model']` and :attr:`config['MODEL_TYPE']`.
If :attr:`config['dataset_save_path']` file exists and
its :attr:`config` of dataset is equal to current :attr:`config` of dataset.
It will return the saved dataset in :attr:`config['dataset_save_path']`.
Args:
config (Config): An instance object of Config, used to record parameter information.
Returns:
Dataset: Constructed dataset.
"""
dataset_module = importlib.import_module('recbole.data.dataset')
if hasattr(dataset_module, config['model'] + 'Dataset'):
dataset_class = getattr(dataset_module, config['model'] + 'Dataset')
else:
model_type = config['MODEL_TYPE']
type2class = {
ModelType.GENERAL: 'Dataset',
ModelType.SEQUENTIAL: 'SequentialDataset',
ModelType.CONTEXT: 'Dataset',
ModelType.KNOWLEDGE: 'KnowledgeBasedDataset',
ModelType.TRADITIONAL: 'Dataset',
ModelType.DECISIONTREE: 'Dataset',
}
dataset_class = getattr(dataset_module, type2class[model_type])
default_file = os.path.join(config['checkpoint_dir'], f'{config["dataset"]}-{dataset_class.__name__}.pth')
file = config['dataset_save_path'] or default_file
if os.path.exists(file):
with open(file, 'rb') as f:
dataset = pickle.load(f)
dataset_args_unchanged = True
for arg in dataset_arguments + ['seed', 'repeatable']:
if config[arg] != dataset.config[arg]:
dataset_args_unchanged = False
break
if dataset_args_unchanged:
logger = getLogger()
logger.info(set_color('Load filtered dataset from', 'pink') + f': [{file}]')
return dataset
dataset = dataset_class(config)
if config['save_dataset']:
dataset.save()
return dataset
def __str__(self):
args_info = '\n'
for category in self.parameters:
args_info += set_color(category + ' Hyper Parameters:\n', 'pink')
args_info += '\n'.join([
(set_color("{}", 'cyan') + " =" +
set_color(" {}", 'yellow')).format(arg, value)
for arg, value in self.final_config_dict.items()
if arg in self.parameters[category]
])
args_info += '\n\n'
args_info += set_color('Other Hyper Parameters: \n', 'pink')
args_info += '\n'.join([
(set_color("{}", 'cyan') + " = " +
set_color("{}", 'yellow')).format(arg, value)
for arg, value in self.final_config_dict.items() if arg not in
{_
for args in self.parameters.values()
for _ in args}.union({'model', 'dataset', 'config_files'})
])
args_info += '\n\n'
return args_info
def data_preparation(config, dataset, save=False):
"""Split the dataset by :attr:`config['eval_args']` and create training, validation and test dataloader.
Args:
config (Config): An instance object of Config, used to record parameter information.
dataset (Dataset): An instance object of Dataset, which contains all interaction records.
save (bool, optional): If ``True``, it will call :func:`save_datasets` to save split dataset.
Defaults to ``False``.
Returns:
tuple:
- train_data (AbstractDataLoader): The dataloader for training.
- valid_data (AbstractDataLoader): The dataloader for validation.
- test_data (AbstractDataLoader): The dataloader for testing.
"""
model_type = config['MODEL_TYPE']
built_datasets = dataset.build()
logger = getLogger()
train_dataset, valid_dataset, test_dataset = built_datasets
train_sampler, valid_sampler, test_sampler = create_samplers(
config, dataset, built_datasets)
if model_type != ModelType.KNOWLEDGE:
train_data = get_dataloader(config, 'train')(config,
train_dataset,
train_sampler,
shuffle=True)
else:
kg_sampler = KGSampler(dataset,
config['train_neg_sample_args']['distribution'])
train_data = get_dataloader(config, 'train')(config,
train_dataset,
train_sampler,
kg_sampler,
shuffle=True)
valid_data = get_dataloader(config, 'evaluation')(config,
valid_dataset,
valid_sampler,
shuffle=False)
test_data = get_dataloader(config, 'evaluation')(config,
test_dataset,
test_sampler,
shuffle=False)
logger.info(
set_color('[Training]: ', 'pink') +
set_color('train_batch_size', 'cyan') + ' = ' +
set_color(f'[{config["train_batch_size"]}]', 'yellow') +
set_color(' negative sampling', 'cyan') + ': ' +
set_color(f'[{config["neg_sampling"]}]', 'yellow'))
logger.info(
set_color('[Evaluation]: ', 'pink') +
set_color('eval_batch_size', 'cyan') + ' = ' +
set_color(f'[{config["eval_batch_size"]}]', 'yellow') +
set_color(' eval_args', 'cyan') + ': ' +
set_color(f'[{config["eval_args"]}]', 'yellow'))
if save:
save_split_dataloaders(config,
dataloaders=(train_data, valid_data, test_data))
return train_data, valid_data, test_data
test_data = get_dataloader(config, 'test')(config,
new_test_dataset,
None,
shuffle=False)
else:
train_data = get_dataloader(config, 'train')(config,
train_dataset,
None,
shuffle=True)
test_data = get_dataloader(config, 'test')(config,
test_dataset,
None,
shuffle=False)
# model loading and initialization
model = get_model(config['model'])(config,
train_data.dataset).to(config['device'])
logger.info(model)
# trainer loading and initialization
trainer = get_trainer(config['MODEL_TYPE'], config['model'])(config, model)
# model training and evaluation
test_score, test_result = trainer.fit(
train_data,
test_data,
saved=True,
show_progress=config['show_progress'])
logger.info(set_color('test result', 'yellow') + f': {test_result}')
def evaluate(self,
eval_data,
load_best_model=True,
model_file=None,
show_progress=False):
r"""Evaluate the model based on the eval data.
Args:
eval_data (DataLoader): the eval data
load_best_model (bool, optional): whether load the best model in the training process, default: True.
It should be set True, if users want to test the model after training.
model_file (str, optional): the saved model file, default: None. If users want to test the previously
trained model file, they can set this parameter.
show_progress (bool): Show the progress of evaluate epoch. Defaults to ``False``.
Returns:
dict: eval result, key is the eval metric and value in the corresponding metric value.
"""
if not eval_data:
return
if load_best_model:
if model_file:
checkpoint_file = model_file
else:
checkpoint_file = self.saved_model_file
checkpoint = torch.load(checkpoint_file)
self.model.load_state_dict(checkpoint['state_dict'])
self.model.load_other_parameter(checkpoint.get('other_parameter'))
message_output = 'Loading model structure and parameters from {}'.format(
checkpoint_file)
self.logger.info(message_output)
self.model.eval()
if isinstance(eval_data, FullSortEvalDataLoader):
eval_func = self._full_sort_batch_eval
if self.item_tensor is None:
self.item_tensor = eval_data.dataset.get_item_feature().to(
self.device)
else:
eval_func = self._neg_sample_batch_eval
if self.config['eval_type'] == EvaluatorType.RANKING:
self.tot_item_num = eval_data.dataset.item_num
iter_data = (tqdm(
eval_data,
total=len(eval_data),
ncols=100,
desc=set_color(f"Evaluate ", 'pink'),
) if show_progress else eval_data)
for batch_idx, batched_data in enumerate(iter_data):
interaction, scores, positive_u, positive_i = eval_func(
batched_data)
if self.gpu_available and show_progress:
iter_data.set_postfix_str(
set_color('GPU RAM: ' + get_gpu_usage(self.device),
'yellow'))
self.eval_collector.eval_batch_collect(scores, interaction,
positive_u, positive_i)
self.eval_collector.model_collect(self.model)
struct = self.eval_collector.get_data_struct()
result = self.evaluator.evaluate(struct)
return result
def fit(self,
train_data,
valid_data=None,
verbose=True,
saved=True,
show_progress=False,
callback_fn=None):
r"""Train the model based on the train data and the valid data.
Args:
train_data (DataLoader): the train data
valid_data (DataLoader, optional): the valid data, default: None.
If it's None, the early_stopping is invalid.
verbose (bool, optional): whether to write training and evaluation information to logger, default: True
saved (bool, optional): whether to save the model parameters, default: True
show_progress (bool): Show the progress of training epoch and evaluate epoch. Defaults to ``False``.
callback_fn (callable): Optional callback function executed at end of epoch.
Includes (epoch_idx, valid_score) input arguments.
Returns:
(float, dict): best valid score and best valid result. If valid_data is None, it returns (-1, None)
"""
if saved and self.start_epoch >= self.epochs:
self._save_checkpoint(-1)
self.eval_collector.data_collect(train_data)
if 'dynamic' in self.config['train_neg_sample_args'].keys(
) and self.config['train_neg_sample_args']['dynamic'] != 'none':
train_data.get_model(self.model)
for epoch_idx in range(self.start_epoch, self.epochs):
# train
training_start_time = time()
train_loss = self._train_epoch(train_data,
epoch_idx,
show_progress=show_progress)
self.train_loss_dict[epoch_idx] = sum(train_loss) if isinstance(
train_loss, tuple) else train_loss
training_end_time = time()
train_loss_output = \
self._generate_train_loss_output(epoch_idx, training_start_time, training_end_time, train_loss)
if verbose:
self.logger.info(train_loss_output)
self._add_train_loss_to_tensorboard(epoch_idx, train_loss)
# eval
if self.eval_step <= 0 or not valid_data:
if saved:
self._save_checkpoint(epoch_idx)
update_output = set_color(
'Saving current',
'blue') + ': %s' % self.saved_model_file
if verbose:
self.logger.info(update_output)
continue
if (epoch_idx + 1) % self.eval_step == 0:
valid_start_time = time()
valid_score, valid_result = self._valid_epoch(
valid_data, show_progress=show_progress)
self.best_valid_score, self.cur_step, stop_flag, update_flag = early_stopping(
valid_score,
self.best_valid_score,
self.cur_step,
max_step=self.stopping_step,
bigger=self.valid_metric_bigger)
valid_end_time = time()
valid_score_output = (set_color("epoch %d evaluating", 'green') + " [" + set_color("time", 'blue')
+ ": %.2fs, " + set_color("valid_score", 'blue') + ": %f]") % \
(epoch_idx, valid_end_time - valid_start_time, valid_score)
valid_result_output = set_color(
'valid result', 'blue') + ': \n' + dict2str(valid_result)
if verbose:
self.logger.info(valid_score_output)
self.logger.info(valid_result_output)
self.tensorboard.add_scalar('Vaild_score', valid_score,
epoch_idx)
if update_flag:
if saved:
self._save_checkpoint(epoch_idx)
update_output = set_color(
'Saving current best',
'blue') + ': %s' % self.saved_model_file
if verbose:
self.logger.info(update_output)
self.best_valid_result = valid_result
if callback_fn:
callback_fn(epoch_idx, valid_score)
if stop_flag:
stop_output = 'Finished training, best eval result in epoch %d' % \
(epoch_idx - self.cur_step * self.eval_step)
if verbose:
self.logger.info(stop_output)
break
self._add_hparam_to_tensorboard(self.best_valid_score)
return self.best_valid_score, self.best_valid_result
