def ner_loss_func(out, *ys, zero=False):
'''
Loss function - to use with fastai learner
It calculates the loss for token classification using softmax cross entropy
If out is already the loss, we simply return the loss
'''
if torch.cuda.is_available():
ys = to_device(ys, torch.cuda.current_device())
# If out is already the loss
if out.size() <= torch.Size([2]):
loss = out.mean() # return mean in case dataparallel is used
else:
loss_fct = torch.nn.CrossEntropyLoss(reduction='none')
if zero:
loss_fct = torch.nn.CrossEntropyLoss(ignore_index=0,
reduction='none')
one, labels, attention_mask = ys
# Only keep active parts of the loss
if attention_mask is not None:
active_loss = attention_mask.view(-1) == 1
active_logits = out.view(-1, len(VOCAB))[active_loss]
active_labels = labels.view(-1)[active_loss]
loss = loss_fct(active_logits, active_labels)
logging.info(active_labels)
logging.info(loss)
logging.info(loss.sum(-1))
loss = loss.mean(-1)
logging.info(loss)
else: # if no attention mask specified calculate loss on all tokens
loss = loss_fct(out.view(-1, len(VOCAB)), labels.view(-1))
return loss
def __next__(self):
if self.idx >= len(self):
raise StopIteration
batch = self.batchdataset[self.idx]
# Note Pinning memory was ~10% _slower_ for the test examples I explored
if self.pin_memory:
batch = _utils.pin_memory.pin_memory_batch(batch)
self.idx = self.idx+1
# move the batch data to device
batch = to_device(batch, self.device)
# return in the form of : xb,yb = (x_cat, x_cont), y
return (batch[0],batch[1]), batch[2]
def on_batch_end(self, last_output, last_target, **kwargs):
pred = last_output.argmax(-1)
true = last_target
if torch.cuda.is_available():
true = to_device(true, torch.cuda.current_device())
_, label_ids, label_mask = true
y_pred = pred.view(-1)
y_true = label_ids.view(-1)
self.predict2 += len(y_pred[y_pred > 1])
preds = y_pred[y_true != 0] # mask of padding
logging.info(y_true)
logging.info(y_pred)
logging.info(preds)
self.predict += len(preds[preds > 1])
self.true += len(y_true[y_true > 1])
self.correct += (np.logical_and(y_true == y_pred,
y_true > 1)).sum().item()
def neighbor_gen(at,
distance_expansion=None,
cutoff=5.0,
n_gaussians=25,
trainable_gaussians=False,
environment_provider=ASEEnvironmentProvider(5.0),
collect_triples=False,
pair_provider=None,
center_positions=True):
properties = {}
properties[Structure.Z] = tensor(at.numbers.astype(np.int)).unsqueeze(0)
positions = at.positions.astype(np.float32)
if center_positions:
positions -= at.get_center_of_mass()
properties[Structure.R] = tensor(positions).unsqueeze(0)
properties[Structure.cell] = tensor(at.cell.astype(
np.float32)).unsqueeze(0)
# get atom environment
idx = 0
nbh_idx, offsets = environment_provider.get_environment(idx, at)
properties[Structure.neighbors] = tensor(nbh_idx.astype(
np.int)).unsqueeze(0)
properties[Structure.cell_offset] = tensor(offsets.astype(
np.float32)).unsqueeze(0)
properties[Structure.neighbor_mask] = None
properties['_idx'] = tensor(np.array([idx], dtype=np.int)).unsqueeze(0)
if collect_triples:
nbh_idx_j, nbh_idx_k = collect_atom_triples(nbh_idx)
properties[Structure.neighbor_pairs_j] = tensor(
nbh_idx_j.astype(np.int))
properties[Structure.neighbor_pairs_k] = tensor(
nbh_idx_k.astype(np.int))
model = spk.custom.representation.RBF(
distance_expansion=distance_expansion,
cutoff=cutoff,
n_gaussians=n_gaussians,
trainable_gaussians=trainable_gaussians)
model = to_device(model)
r, f = model.forward(properties)
return to_np(r.squeeze()), to_np(f.squeeze())
def __next__(self):
if self.idx >= len(self):
raise StopIteration
if self.batchdataset.mode == 'test':
X, y, infor = self.batchdataset[self.idx]
batch = (X, y)
else:
batch = self.batchdataset[self.idx]
# Note Pinning memory was ~10% _slower_ for the test examples I explored
if self.pin_memory:
batch = _utils.pin_memory.pin_memory_batch(batch)
self.idx = self.idx + 1
# move the batch data to device
batch = to_device(batch, self.device)
# return in the form of : xb,yb = (x_cat, x_cont), y
if self.batchdataset.mode == 'test':
return batch, infor
return batch
def replace_bn(learner, m, G=32):
for n, ch in m.named_children():
if isinstance(ch, nn.BatchNorm2d):
new_layer = to_device(
nn.GroupNorm(G, ch.num_features, ch.eps, ch.affine),
learner.data.device)
m._modules[n] = new_layer
found = False
for lg in learner.layer_groups:
for k, c in lg.named_children():
if c is ch:
lg._modules[k] = new_layer
found = True
break
if found:
break
for ch in m.children():
replace_bn(learner, ch)
if hasattr(learner, 'opt'):
learner.create_opt(learner.opt.lr, wd=learner.wd)
def conll_f1(pred, *true, eps: float = 1e-9):
''' NOTE: calulcates F1 per batch
- use Conll_F1 callback class to calculate overall F1 score
'''
if torch.cuda.is_available():
true = to_device(true, torch.cuda.current_device())
pred = pred.argmax(-1)
_, label_ids, label_mask = true
mask = label_mask.view(-1) == 1
y_pred = pred.view(-1)[mask]
y_true = label_ids.view(-1)[mask]
all_pos = len(y_pred[y_pred > 1])
actual_pos = len(y_true[y_true > 1])
correct_pos = (np.logical_and(y_true == y_pred, y_true > 1)).sum().item()
logging.info(f'{all_pos} - {actual_pos} -> {correct_pos}')
prec = correct_pos / (all_pos + eps)
rec = correct_pos / (actual_pos + eps)
f1 = (2 * prec * rec) / (prec + rec + eps)
logging.info(f'f1: {f1} prec: {prec}, rec: {rec}')
return torch.Tensor([f1])
def unet_learner_wide(
data: DataBunch,
arch: Callable,
pretrained: bool = True,
blur_final: bool = True,
norm_type: Optional[NormType] = NormType,
split_on: Optional[SplitFuncOrIdxList] = None,
blur: bool = False,
self_attention: bool = False,
y_range: Optional[Tuple[float, float]] = None,
last_cross: bool = True,
bottle: bool = False,
nf_factor: int = 1,
**kwargs: Any
) -> Learner:
"Build Unet learner from `data` and `arch`."
meta = cnn_config(arch)
body = create_body(arch, pretrained)
model = to_device(
DynamicUnetWide(
body,
n_classes=data.c,
blur=blur,
blur_final=blur_final,
self_attention=self_attention,
y_range=y_range,
norm_type=norm_type,
last_cross=last_cross,
bottle=bottle,
nf_factor=nf_factor,
),
data.device,
)
learn = Learner(data, model, **kwargs)
learn.split(ifnone(split_on, meta['split']))
if pretrained:
learn.freeze()
apply_init(model[2], nn.init.kaiming_normal_)
return learn
def unet_learner_without_skip_connections(n_classes, device, arch:Callable, pretrained:bool=True, blur_final:bool=True,
norm_type:Optional[NormType]=NormType, split_on:Optional[SplitFuncOrIdxList]=None, blur:bool=False,
y_range:Optional[Tuple[float,float]]=None, skip_connections=True,
cut:Union[int,Callable]=None, **learn_kwargs:Any):
"Build Unet learner from `data` and `arch`."
from fastai.vision import create_body
from fastai.torch_core import to_device
from fastai.torch_core import apply_init
# meta = cnn_config(arch)
body = create_body(arch, pretrained, cut)
# noinspection PyTypeChecker
model = to_device(DynamicUnetWithoutSkipConnections(body, n_classes=n_classes, y_range=y_range, norm_type=norm_type,
skip_connections=skip_connections),
device)
# learn = Learner(data, model, **learn_kwargs)
# learn.split(ifnone(split_on, meta['split']))
# if pretrained: learn.freeze()
apply_init(model[2], nn.init.kaiming_normal_)
return model
def __call__(self, b, tfm_y=TfmY.No):
return to_device(b, self.device)
