[FEATURE] Remove utility modules and add wrappers instead

This commit is contained in:
Jensun Ravichandran 2021-06-04 22:16:55 +02:00
parent 7bc0bfa3ab
commit 0c28eda706
2 changed files with 36 additions and 16 deletions

View File

@ -1,16 +0,0 @@
"""ProtoTorch utility modules."""
import torch
class LambdaLayer(torch.nn.Module):
def __init__(self, fn, name=None):
super().__init__()
self.fn = fn
self.name = name or fn.__name__ # lambda fns get <lambda>
def forward(self, *args, **kwargs):
return self.fn(*args, **kwargs)
def extra_repr(self):
return self.name

View File

@ -0,0 +1,36 @@
"""ProtoTorch Wrappers."""
import torch
class LambdaLayer(torch.nn.Module):
def __init__(self, fn, name=None):
super().__init__()
self.fn = fn
self.name = name or fn.__name__ # lambda fns get <lambda>
def forward(self, *args, **kwargs):
return self.fn(*args, **kwargs)
def extra_repr(self):
return self.name
class LossLayer(torch.nn.modules.loss._Loss):
def __init__(self,
fn,
name=None,
size_average=None,
reduce=None,
reduction: str = "mean") -> None:
super().__init__(size_average=size_average,
reduce=reduce,
reduction=reduction)
self.fn = fn
self.name = name or fn.__name__ # lambda fns get <lambda>
def forward(self, *args, **kwargs):
return self.fn(*args, **kwargs)
def extra_repr(self):
return self.name