-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathModelClass.py
128 lines (102 loc) · 4.49 KB
/
ModelClass.py
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import torch
import torch.nn as nn
class Swish(nn.Module):
def __init__(self, ):
super().__init__()
def forward(self, x):
return x * torch.sigmoid(x)
class Gaussian(nn.Module):
def __init__(self, ):
super().__init__()
def forward(self, x):
return torch.exp(-5 * x ** 2)
class Sin(nn.Module):
def __init__(self, ):
super().__init__()
def forward(self, x):
return torch.sin(x)
class Snake(nn.Module):
def __init__(self):
super().__init__()
self.alpha = 0.5
def forward(self, x):
return x + torch.sin(self.alpha * x) ** 2 / self.alpha
def activation(name):
if name in ['tanh', 'Tanh']:
return nn.Tanh()
elif name in ['relu', 'ReLU']:
return nn.ReLU(inplace=True)
elif name in ['lrelu', 'LReLU']:
return nn.LeakyReLU(inplace=True)
elif name in ['sigmoid', 'Sigmoid']:
return nn.Sigmoid()
elif name in ['softplus', 'Softplus']:
return nn.Softplus(beta=4)
elif name in ['celu', 'CeLU']:
return nn.CELU()
elif name in ['sin', 'Sin']:
return Sin()
elif name in ['swish']:
return Swish()
elif name in ['snake']:
return Snake()
elif name in ['gaussian']:
return Gaussian()
else:
raise ValueError('Unknown activation function')
class Pinns(nn.Module):
def __init__(self, input_dimension, output_dimension, network_properties):
super(Pinns, self).__init__()
self.input_dimension = input_dimension
self.output_dimension = output_dimension
self.n_hidden_layers = int(network_properties["hidden_layers"])
self.neurons = int(network_properties["neurons"])
self.lambda_residual = float(network_properties["residual_parameter"])
self.kernel_regularizer = int(network_properties["kernel_regularizer"])
self.regularization_param = float(network_properties["regularization_parameter"])
self.num_epochs = int(network_properties["epochs"])
self.act_string = str(network_properties["activation"])
self.iterations = int(network_properties["iterations"])
self.reset_freq = network_properties["reset_freq"]
self.loss = network_properties["loss_type"]
self.activation = activation(self.act_string)
if self.n_hidden_layers != 0:
self.input_layer = nn.Linear(self.input_dimension, self.neurons)
self.hidden_layers = nn.ModuleList([nn.Linear(self.neurons, self.neurons) for _ in range(self.n_hidden_layers - 1)])
self.output_layer = nn.Linear(self.neurons, self.output_dimension)
else:
self.input_output_layer = nn.Linear(self.input_dimension, self.output_dimension)
def forward(self, x):
if self.n_hidden_layers != 0:
x = self.activation(self.input_layer(x))
for k, l in enumerate(self.hidden_layers):
x = self.activation(l(x))
return self.output_layer(x)
else:
return self.input_output_layer(x)
class PinnsTest(nn.Module):
def __init__(self, input_dimension, output_dimension, network_properties):
super(PinnsTest, self).__init__()
self.input_dimension = input_dimension
self.output_dimension = output_dimension
self.n_hidden_layers = int(network_properties["hidden_layers"])
self.neurons = int(network_properties["neurons"])
self.lambda_residual = float(network_properties["residual_parameter"])
self.kernel_regularizer = int(network_properties["kernel_regularizer"])
self.regularization_param = float(network_properties["regularization_parameter"])
self.num_epochs = int(network_properties["epochs"])
self.act_string = str(network_properties["activation"])
self.iterations = int(network_properties["iterations"])
self.reset_freq = (network_properties["reset_freq"])
self.input_layer = nn.Linear(self.input_dimension, self.neurons)
self.hidden_layers = nn.ModuleList([nn.Linear(self.neurons, self.neurons) for _ in range(self.n_hidden_layers - 1)])
self.output_layer = nn.Linear(self.neurons, self.output_dimension)
self.activation = activation(self.act_string)
self.activation_out = Gaussian()
self.activation_out_2 = torch.nn.Softplus()
def forward(self, x):
x = self.activation(self.input_layer(x))
for k, l in enumerate(self.hidden_layers):
x = self.activation(l(x))
o = self.output_layer(x)
return o