[FEATURE] Remove utility modules and add wrappers instead
This commit is contained in:
parent
7bc0bfa3ab
commit
0c28eda706
@ -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
|
|
36
prototorch/modules/wrappers.py
Normal file
36
prototorch/modules/wrappers.py
Normal 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
|
Loading…
Reference in New Issue
Block a user