def _train(self, loader):
self.model.train()
bar, metric = progress_bar(loader), ChartMetric()
for i, (words, *feats, labels) in enumerate(bar, 1):
word_mask = words.ne(self.args.pad_index)
mask = word_mask if len(words.shape) < 3 else word_mask.any(-1)
mask = mask.unsqueeze(1) & mask.unsqueeze(2)
mask[:, 0] = 0
s_edge, s_sib, s_cop, s_grd, s_label = self.model(words, feats)
loss, s_edge = self.model.loss(s_edge, s_sib, s_cop, s_grd,
s_label, labels, mask)
loss = loss / self.args.update_steps
loss.backward()
nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
if i % self.args.update_steps == 0:
self.optimizer.step()
self.scheduler.step()
self.optimizer.zero_grad()
label_preds = self.model.decode(s_edge, s_label)
metric(label_preds.masked_fill(~mask, -1),
labels.masked_fill(~mask, -1))
bar.set_postfix_str(
f"lr: {self.scheduler.get_last_lr()[0]:.4e} - loss: {loss:.4f} - {metric}"
)
logger.info(f"{bar.postfix}")
def _train(self, loader):
self.model.train()
bar, metric = progress_bar(loader), ChartMetric()
for words, *feats, edges, labels in bar:
self.optimizer.zero_grad()
mask = words.ne(self.WORD.pad_index)
mask = mask.unsqueeze(1) & mask.unsqueeze(2)
mask[:, 0] = 0
s_edge, s_sib, s_cop, s_grd, s_label = self.model(words, feats)
loss, s_edge = self.model.loss(s_edge, s_sib, s_cop, s_grd,
s_label, edges, labels, mask)
loss.backward()
nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
self.optimizer.step()
self.scheduler.step()
edge_preds, label_preds = self.model.decode(s_edge, s_label)
metric(label_preds.masked_fill(~(edge_preds.gt(0) & mask), -1),
labels.masked_fill(~(edges.gt(0) & mask), -1))
bar.set_postfix_str(
f"lr: {self.scheduler.get_last_lr()[0]:.4e} - loss: {loss:.4f} - {metric}"
)
def _evaluate(self, loader):
self.model.eval()
total_loss, metric = 0, ChartMetric()
for words, feats, edges, labels in loader:
mask = words.ne(self.WORD.pad_index)
mask = mask.unsqueeze(1) & mask.unsqueeze(2)
s_edge, s_label = self.model(words, feats)
loss = self.model.loss(s_edge, s_label, edges, labels, mask)
total_loss += loss.item()
chart_preds = self.model.decode(s_edge, s_label, mask)
metric(chart_preds, labels.masked_fill(~(edges.gt(0) & mask), -1))
total_loss /= len(loader)
return total_loss, metric
def _evaluate(self, loader):
self.model.eval()
total_loss, metric = 0, ChartMetric()
for words, *feats, labels in loader:
word_mask = words.ne(self.args.pad_index)
mask = word_mask if len(words.shape) < 3 else word_mask.any(-1)
mask = mask.unsqueeze(1) & mask.unsqueeze(2)
mask[:, 0] = 0
s_edge, s_label = self.model(words, feats)
loss = self.model.loss(s_edge, s_label, labels, mask)
total_loss += loss.item()
label_preds = self.model.decode(s_edge, s_label)
metric(label_preds.masked_fill(~mask, -1),
labels.masked_fill(~mask, -1))
total_loss /= len(loader)
return total_loss, metric
