julius/prototorch_models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[REFACTOR] Use LambdaLayer instead of EuclideanDistance

Browse Source
This commit is contained in:
Jensun Ravichandran 2021-06-02 00:21:11 +02:00
parent ef4d70eee0
commit 98c198d463
1 changed files with 3 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
prototorch/models/unsupervised.py
View File

@ -10,17 +10,13 @@ from prototorch.components import Components, LabeledComponents
from prototorch.components.initializers import ZerosInitializer, parse_data_arg
from prototorch.functions.competitions import knnc
from prototorch.functions.distances import euclidean_distance
from prototorch.modules import LambdaLayer
from prototorch.modules.losses import NeuralGasEnergy
from pytorch_lightning.callbacks import Callback
from .abstract import AbstractPrototypeModel
class EuclideanDistance(torch.nn.Module):
def forward(self, x, y):
return euclidean_distance(x, y)
class GNGCallback(Callback):
"""GNG Callback.
@ -201,7 +197,7 @@ class NeuralGas(AbstractPrototypeModel):
self.hparams.num_prototypes,
initializer=self.hparams.prototype_initializer)
self.distance_layer = EuclideanDistance()
self.distance_layer = LambdaLayer(euclidean_distance)
self.energy_layer = NeuralGasEnergy(lm=self.hparams.lm)
self.topology_layer = ConnectionTopology(
agelimit=self.hparams.agelimit,
@ -212,8 +208,7 @@ class NeuralGas(AbstractPrototypeModel):
x = train_batch[0]
protos = self.proto_layer()
d = self.distance_layer(x, protos)
cost, order = self.energy_layer(d)
cost, _ = self.energy_layer(d)
self.topology_layer(d)
return cost
@ -235,9 +230,7 @@ class GrowingNeuralGas(NeuralGas):
protos = self.proto_layer()
d = self.distance_layer(x, protos)
cost, order = self.energy_layer(d)
winner = order[:, 0]
mask = torch.zeros_like(d)
mask[torch.arange(len(mask)), winner] = 1.0
winner_distances = d * mask