[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