File size: 1,868 Bytes
4f5540c |
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 |
import torch
class Vector_IV(torch.nn.Module):
def __init__(self, input_feat, hidden_channels):
super(Vector_IV, self).__init__()
self.fc = torch.nn.Sequential(
torch.nn.Linear(input_feat, hidden_channels),
torch.nn.PReLU(),
torch.nn.Linear(hidden_channels, 1)
)
def forward(self, x):
return self.fc(x)
class Vector_Tg(torch.nn.Module):
def __init__(self, input_feat, hidden_channels):
super(Vector_Tg, self).__init__()
self.fc = torch.nn.Sequential(
torch.nn.Linear(input_feat, hidden_channels),
torch.nn.PReLU(),
)
self.TG = torch.nn.Linear(hidden_channels, 1)
self.mult_factor = torch.nn.Sequential(
torch.nn.Linear(hidden_channels, 1),
torch.nn.Tanh(),
)
def forward(self, x):
x = self.fc(x)
return self.TG(x) * self.mult_factor(x)
class Vector_Joint(torch.nn.Module):
def __init__(self, input_feat, hidden_channels):
super(Vector_Joint, self).__init__()
self.fc1 = torch.nn.Linear(input_feat, hidden_channels)
self.leaky1 = torch.nn.PReLU()
self.fc2 = torch.nn.Sequential( # IV model
torch.nn.Linear(hidden_channels, hidden_channels),
torch.nn.PReLU(),
torch.nn.Linear(hidden_channels, 1)
)
self.fc3 = torch.nn.Linear(hidden_channels, 1) # Tg
self.mult_factor = torch.nn.Sequential(
torch.nn.Linear(hidden_channels, 1),
torch.nn.Tanh(),
)
def forward(self, x):
x = self.leaky1(self.fc1(x))
# Run IV:
x_IV = self.fc2(x)
# Run Tg (+ multiplying factor):
x_Tg = torch.exp(self.fc3(x))
m = self.mult_factor(x)
x_Tg = x_Tg * m
return {'IV': x_IV, 'Tg': x_Tg} |