def predict_level(self, level, test_x, k, labels_num):
data_cnf, model_cnf = self.data_cnf, self.model_cnf
model = self.models.get(level, None)
if level == 0:
logger.info(F'Predicting Level-{level}, Top: {k}')
if model is None:
model = Model(AttentionRNN, labels_num=labels_num, model_path=F'{self.model_path}-Level-{level}',
emb_init=self.emb_init, **data_cnf['model'], **model_cnf['model'])
test_loader = DataLoader(MultiLabelDataset(test_x), model_cnf['predict']['batch_size'],
num_workers=4)
return model.predict(test_loader, k=k)
else:
if level == self.level - 1:
groups = np.load(F'{self.groups_path}-Level-{level-1}.npy')
else:
groups = self.get_inter_groups(labels_num)
group_scores, group_labels = self.predict_level(level - 1, test_x, self.top, len(groups))
torch.cuda.empty_cache()
logger.info(F'Predicting Level-{level}, Top: {k}')
if model is None:
model = XMLModel(network=FastAttentionRNN, labels_num=labels_num,
model_path=F'{self.model_path}-Level-{level}',
emb_init=self.emb_init, **data_cnf['model'], **model_cnf['model'])
test_loader = DataLoader(XMLDataset(test_x, labels_num=labels_num,
groups=groups, group_labels=group_labels, group_scores=group_scores),
model_cnf['predict']['batch_size'], num_workers=4)
return model.predict(test_loader, k=k)
def train_level(self, level, train_x, train_y, valid_x, valid_y):
model_cnf, data_cnf = self.model_cnf, self.data_cnf
if level == 0:
while not os.path.exists(F'{self.groups_path}-Level-{level}.npy'):
time.sleep(30)
groups = np.load(F'{self.groups_path}-Level-{level}.npy')
train_y, valid_y = self.get_mapping_y(groups, self.labels_num,
train_y, valid_y)
labels_num = len(groups)
train_loader = DataLoader(MultiLabelDataset(train_x, train_y),
model_cnf['train'][level]['batch_size'],
num_workers=4,
shuffle=True)
valid_loader = DataLoader(MultiLabelDataset(valid_x,
valid_y,
training=False),
model_cnf['valid']['batch_size'],
num_workers=4)
model = Model(AttentionRNN,
labels_num=labels_num,
model_path=F'{self.model_path}-Level-{level}',
emb_init=self.emb_init,
**data_cnf['model'],
**model_cnf['model'])
if not os.path.exists(model.model_path):
logger.info(
F'Training Level-{level}, Number of Labels: {labels_num}')
model.train(train_loader, valid_loader,
**model_cnf['train'][level])
model.optimizer = None
logger.info(F'Finish Training Level-{level}')
self.models[level] = model
logger.info(F'Generating Candidates for Level-{level+1}, '
F'Number of Labels: {labels_num}, Top: {self.top}')
train_loader = DataLoader(MultiLabelDataset(train_x),
model_cnf['valid']['batch_size'],
num_workers=4)
return train_y, model.predict(
train_loader, k=self.top), model.predict(valid_loader,
k=self.top)
else:
train_group_y, train_group, valid_group = self.train_level(
level - 1, train_x, train_y, valid_x, valid_y)
torch.cuda.empty_cache()
logger.info('Getting Candidates')
_, group_labels = train_group
group_candidates = np.empty((len(train_x), self.top), dtype=np.int)
for i, labels in tqdm(enumerate(group_labels),
leave=False,
desc='Parents'):
ys, ye = train_group_y.indptr[i], train_group_y.indptr[i + 1]
positive = set(train_group_y.indices[ys:ye])
if self.top >= len(positive):
candidates = positive
for la in labels:
if len(candidates) == self.top:
break
if la not in candidates:
candidates.add(la)
else:
candidates = set()
for la in labels:
if la in positive:
candidates.add(la)
if len(candidates) == self.top:
break
if len(candidates) < self.top:
candidates = (list(candidates) +
list(positive - candidates))[:self.top]
group_candidates[i] = np.asarray(list(candidates))
if level < self.level - 1:
while not os.path.exists(
F'{self.groups_path}-Level-{level}.npy'):
time.sleep(30)
groups = np.load(F'{self.groups_path}-Level-{level}.npy')
train_y, valid_y = self.get_mapping_y(groups, self.labels_num,
train_y, valid_y)
labels_num, last_groups = len(groups), self.get_inter_groups(
len(groups))
else:
groups, labels_num = None, train_y.shape[1]
last_groups = np.load(
F'{self.groups_path}-Level-{level-1}.npy')
train_loader = DataLoader(XMLDataset(
train_x,
train_y,
labels_num=labels_num,
groups=last_groups,
group_labels=group_candidates),
model_cnf['train'][level]['batch_size'],
num_workers=4,
shuffle=True)
group_scores, group_labels = valid_group
valid_loader = DataLoader(XMLDataset(valid_x,
valid_y,
training=False,
labels_num=labels_num,
groups=last_groups,
group_labels=group_labels,
group_scores=group_scores),
model_cnf['valid']['batch_size'],
num_workers=4)
model = XMLModel(network=FastAttentionRNN,
labels_num=labels_num,
emb_init=self.emb_init,
model_path=F'{self.model_path}-Level-{level}',
**data_cnf['model'],
**model_cnf['model'])
if not os.path.exists(model.model_path):
logger.info(
F'Loading parameters of Level-{level} from Level-{level-1}'
)
last_model = self.get_last_models(level - 1)
model.network.module.emb.load_state_dict(
last_model.module.emb.state_dict())
model.network.module.lstm.load_state_dict(
last_model.module.lstm.state_dict())
model.network.module.linear.load_state_dict(
last_model.module.linear.state_dict())
logger.info(
F'Training Level-{level}, '
F'Number of Labels: {labels_num}, '
F'Candidates Number: {train_loader.dataset.candidates_num}'
)
model.train(train_loader, valid_loader,
**model_cnf['train'][level])
model.optimizer = model.state = None
logger.info(F'Finish Training Level-{level}')
self.models[level] = model
if level == self.level - 1:
return
logger.info(F'Generating Candidates for Level-{level+1}, '
F'Number of Labels: {labels_num}, Top: {self.top}')
group_scores, group_labels = train_group
train_loader = DataLoader(XMLDataset(train_x,
labels_num=labels_num,
groups=last_groups,
group_labels=group_labels,
group_scores=group_scores),
model_cnf['valid']['batch_size'],
num_workers=4)
return train_y, model.predict(
train_loader, k=self.top), model.predict(valid_loader,
k=self.top)