wehub-resource-sync
150 lines
4.7 KiB
Python
150 lines
4.7 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.distributed.fleet import auto
|
|
from paddle.incubate.autograd import Hessian
|
|
|
|
np.random.seed(1234)
|
|
paddle.seed(1234)
|
|
|
|
|
|
class FCNet:
|
|
def __init__(self, num_ins, num_outs, num_layers, hidden_size):
|
|
self.num_ins = num_ins
|
|
self.num_outs = num_outs
|
|
self.num_layers = num_layers
|
|
self.hidden_size = hidden_size
|
|
self.activation = paddle.tanh
|
|
|
|
self.weights = []
|
|
self.biases = []
|
|
for i in range(self.num_layers):
|
|
if i == 0:
|
|
lsize = self.num_ins
|
|
rsize = self.hidden_size
|
|
elif i == (self.num_layers - 1):
|
|
lsize = self.hidden_size
|
|
rsize = self.num_outs
|
|
else:
|
|
lsize = self.hidden_size
|
|
rsize = self.hidden_size
|
|
|
|
w = paddle.static.create_parameter(
|
|
shape=[lsize, rsize], dtype="float32", is_bias=False
|
|
)
|
|
b = paddle.static.create_parameter(
|
|
shape=[rsize], dtype="float32", is_bias=True
|
|
)
|
|
self.weights.append(w)
|
|
self.biases.append(b)
|
|
|
|
def nn_func(self, ins):
|
|
u = ins
|
|
for i in range(self.num_layers - 1):
|
|
u = paddle.nn.functional.linear(u, self.weights[i], self.biases[i])
|
|
u = self.activation(u)
|
|
u = paddle.nn.functional.linear(u, self.weights[-1], self.biases[-1])
|
|
return u
|
|
|
|
|
|
class LaplaceModel(paddle.nn.Layer):
|
|
def __init__(self, num_ins=2, num_outs=1, num_layers=5, hidden_size=20):
|
|
super().__init__()
|
|
self.net = FCNet(
|
|
num_ins=num_ins,
|
|
num_outs=num_outs,
|
|
num_layers=num_layers,
|
|
hidden_size=hidden_size,
|
|
)
|
|
|
|
def forward(self, inputs, bc_index):
|
|
inputs.stop_gradient = False
|
|
outputs = self.net.nn_func(inputs)
|
|
# eq_loss
|
|
hes = Hessian(self.net.nn_func, inputs, is_batched=True)
|
|
eq_loss = paddle.norm(hes[:, 0, 0] + hes[:, 1, 1], p=2)
|
|
# bc_loss
|
|
bc_u = paddle.index_select(outputs, bc_index)
|
|
return eq_loss, bc_u
|
|
|
|
|
|
class LaplaceDataset(paddle.io.Dataset):
|
|
def __init__(self, num_sample):
|
|
self.num_sample = num_sample
|
|
|
|
def __getitem__(self, index):
|
|
x = np.linspace(0, 0.9, 10)
|
|
y = np.linspace(0, 0.9, 10)
|
|
np.random.seed(index) # Optional: Ensure reproducibility
|
|
bc_value = np.random.rand(36).reshape(36, 1).astype('float32')
|
|
|
|
domain_space = []
|
|
bc_index = []
|
|
for j in range(len(y)):
|
|
for i in range(len(x)):
|
|
domain_space.append([x[i], y[j]])
|
|
if i == 0 or i == 9 or j == 0 or j == 9:
|
|
bc_index.append(i + 10 * j)
|
|
domain_space = np.array(domain_space, dtype='float32')
|
|
bc_index = np.array(bc_index, dtype='int64')
|
|
# Return a single input point and its related information based on the index
|
|
idx = index % len(domain_space)
|
|
return domain_space[idx], bc_index, bc_value
|
|
|
|
def __len__(self):
|
|
return self.num_sample
|
|
|
|
|
|
def loss_func(eq_loss, bc_u, bc_value):
|
|
bc_diff = bc_u - bc_value
|
|
bc_loss = paddle.norm(bc_diff, p=2)
|
|
loss = eq_loss + bc_loss
|
|
return loss
|
|
|
|
|
|
def main():
|
|
paddle.enable_static()
|
|
# dataset
|
|
train_dataset = LaplaceDataset(10)
|
|
# optimizer
|
|
optimizer = paddle.optimizer.Adam(learning_rate=0.001)
|
|
# model
|
|
laplace = LaplaceModel()
|
|
|
|
dist_strategy = auto.Strategy()
|
|
dist_strategy.auto_mode = "semi"
|
|
|
|
engine = auto.Engine(
|
|
laplace, loss=loss_func, optimizer=optimizer, strategy=dist_strategy
|
|
)
|
|
engine.fit(train_dataset, train_sample_split=2, batch_size=None)
|
|
|
|
dist_context = engine.dist_context
|
|
block = engine.main_program.global_block()
|
|
ops = block.ops
|
|
for op in ops:
|
|
if op.type == 'p_norm':
|
|
op_dist_attr = dist_context.get_op_dist_attr_for_program(op)
|
|
assert op_dist_attr.impl_type == 'p_norm'
|
|
if 'x' in op.input_arg_names:
|
|
out_name = op.output_arg_names[0]
|
|
assert block.vars[out_name].shape[0] == 50
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|