File size: 3,823 Bytes
7b5e67a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
from torch import nn
class SingleVisualizationModel(nn.Module):
def __init__(self, input_dims, output_dims, units, hidden_layer=3):
super(SingleVisualizationModel, self).__init__()
self.input_dims = input_dims
self.output_dims = output_dims
self.units = units
self.hidden_layer = hidden_layer
self._init_autoencoder()
# TODO find the best model architecture
def _init_autoencoder(self):
self.encoder = nn.Sequential(
nn.Linear(self.input_dims, self.units),
nn.ReLU(True))
for h in range(self.hidden_layer):
self.encoder.add_module("{}".format(2*h+2), nn.Linear(self.units, self.units))
self.encoder.add_module("{}".format(2*h+3), nn.ReLU(True))
self.encoder.add_module("{}".format(2*(self.hidden_layer+1)), nn.Linear(self.units, self.output_dims))
self.decoder = nn.Sequential(
nn.Linear(self.output_dims, self.units),
nn.ReLU(True))
for h in range(self.hidden_layer):
self.decoder.add_module("{}".format(2*h+2), nn.Linear(self.units, self.units))
self.decoder.add_module("{}".format(2*h+3), nn.ReLU(True))
self.decoder.add_module("{}".format(2*(self.hidden_layer+1)), nn.Linear(self.units, self.input_dims))
def forward(self, edge_to, edge_from):
outputs = dict()
embedding_to = self.encoder(edge_to)
embedding_from = self.encoder(edge_from)
recon_to = self.decoder(embedding_to)
recon_from = self.decoder(embedding_from)
outputs["umap"] = (embedding_to, embedding_from)
outputs["recon"] = (recon_to, recon_from)
return outputs
class VisModel(nn.Module):
"""define you own visualizatio model by specifying the structure
"""
def __init__(self, encoder_dims, decoder_dims):
"""define you own visualizatio model by specifying the structure
Parameters
----------
encoder_dims : list of int
the neuron number of your encoder
for example, [100,50,2], denote two fully connect layers, with shape (100,50) and (50,2)
decoder_dims : list of int
same as encoder_dims
"""
super(VisModel, self).__init__()
assert len(encoder_dims) > 1
assert len(decoder_dims) > 1
self.encoder_dims = encoder_dims
self.decoder_dims = decoder_dims
self._init_autoencoder()
def _init_autoencoder(self):
self.encoder = nn.Sequential()
for i in range(0, len(self.encoder_dims)-2):
self.encoder.add_module("{}".format(len(self.encoder)), nn.Linear(self.encoder_dims[i], self.encoder_dims[i+1]))
self.encoder.add_module("{}".format(len(self.encoder)), nn.ReLU(True))
self.encoder.add_module("{}".format(len(self.encoder)), nn.Linear(self.encoder_dims[-2], self.encoder_dims[-1]))
self.decoder = nn.Sequential()
for i in range(0, len(self.decoder_dims)-2):
self.decoder.add_module("{}".format(len(self.decoder)), nn.Linear(self.decoder_dims[i], self.decoder_dims[i+1]))
self.decoder.add_module("{}".format(len(self.decoder)), nn.ReLU(True))
self.decoder.add_module("{}".format(len(self.decoder)), nn.Linear(self.decoder_dims[-2], self.decoder_dims[-1]))
def forward(self, edge_to, edge_from):
outputs = dict()
embedding_to = self.encoder(edge_to)
embedding_from = self.encoder(edge_from)
recon_to = self.decoder(embedding_to)
recon_from = self.decoder(embedding_from)
outputs["umap"] = (embedding_to, embedding_from)
outputs["recon"] = (recon_to, recon_from)
return outputs
'''
The visualization model definition class
'''
|