wangmin
5 years ago
committed by
mwang
parent
b5aec1fd45
commit
7e0abae743
@ -0,0 +1,42 @@
|
||||
#!/bin/bash
|
||||
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# 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.
|
||||
# ============================================================================
|
||||
|
||||
if [ $# != 5 ]
|
||||
then
|
||||
echo "Usage: sh run_standalone_train_gpu.sh [im] [jm] [kb] [step] [DATASET_PATH]"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
|
||||
ulimit -u unlimited
|
||||
export DEVICE_NUM=1
|
||||
export DEVICE_ID=0
|
||||
|
||||
if [ -d "train" ];
|
||||
then
|
||||
rm -rf ./train
|
||||
fi
|
||||
mkdir ./train
|
||||
cp ../*.py ./train
|
||||
cp *.sh ./train
|
||||
cp -r ../src ./train
|
||||
cd ./train || exit
|
||||
mkdir ./outputs
|
||||
echo "start training for device $DEVICE_ID"
|
||||
env > env.log
|
||||
python train.py --im=$1 --jm=$2 --kb=$3 --step=$4 --file_path=$5 --outputs_path="./outputs/" &> log &
|
||||
|
||||
cd ..
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,57 @@
|
||||
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# 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.
|
||||
# ============================================================================
|
||||
"""Grid"""
|
||||
|
||||
import mindspore.nn as nn
|
||||
from src.stencil import AXB, AYB, AZB
|
||||
|
||||
|
||||
class Grid(nn.Cell):
|
||||
"""
|
||||
init C grid
|
||||
"""
|
||||
def __init__(self, im, jm, km, stencil_width=1):
|
||||
super(Grid, self).__init__()
|
||||
self.im = im
|
||||
self.jm = jm
|
||||
self.km = km
|
||||
self.x_map = [1, 0, 3, 2, 5, 4, 7, 6]
|
||||
self.y_map = [2, 3, 0, 1, 6, 7, 4, 5]
|
||||
self.z_map = [4, 5, 6, 7, 0, 1, 2, 3]
|
||||
self.AXB = AXB(stencil_width=stencil_width)
|
||||
self.AYB = AYB(stencil_width=stencil_width)
|
||||
self.AZB = AZB(stencil_width=stencil_width)
|
||||
|
||||
def construct(self, dx, dy, dz):
|
||||
"""construct"""
|
||||
dx0 = self.AYB(self.AXB(dx))
|
||||
dy0 = self.AYB(self.AXB(dy))
|
||||
dz0 = dz
|
||||
|
||||
dx1 = self.AYB(dx)
|
||||
dy1 = self.AYB(dy)
|
||||
dz1 = self.AZB(dz)
|
||||
|
||||
dx2 = self.AXB(dx)
|
||||
dy2 = self.AXB(dy)
|
||||
|
||||
dx3 = dx
|
||||
dy3 = dy
|
||||
|
||||
x_d = (dx0, dx1, dx2, dx3)
|
||||
y_d = (dy0, dy1, dy2, dy3)
|
||||
z_d = (dz0, dz1)
|
||||
|
||||
return x_d, y_d, z_d
|
||||
@ -0,0 +1,210 @@
|
||||
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# 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.
|
||||
# ============================================================================
|
||||
"""stencil operations kernel"""
|
||||
|
||||
import mindspore.nn as nn
|
||||
from mindspore.ops import operations as P
|
||||
|
||||
|
||||
class axb_kernel(nn.Cell):
|
||||
"""create axb_kernel"""
|
||||
def __init__(self):
|
||||
super(axb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((1, 0), (0, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class ayb_kernel(nn.Cell):
|
||||
"""create ayb_kernel"""
|
||||
def __init__(self):
|
||||
super(ayb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (1, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class azb_kernel(nn.Cell):
|
||||
"""create azb_kernel"""
|
||||
def __init__(self):
|
||||
super(azb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 0), (1, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class axf_kernel(nn.Cell):
|
||||
"""create axf_kernel"""
|
||||
def __init__(self):
|
||||
super(axf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 1), (0, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (1, 0, 0), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class ayf_kernel(nn.Cell):
|
||||
"""create ayf_kernel"""
|
||||
def __init__(self):
|
||||
super(ayf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 1), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 1, 0), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class azf_kernel(nn.Cell):
|
||||
"""create azf_kernel"""
|
||||
def __init__(self):
|
||||
super(azf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 0), (0, 1)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 1), x_shape)
|
||||
out = 0.5 * (x + x1)
|
||||
return out
|
||||
|
||||
|
||||
class dxb_kernel(nn.Cell):
|
||||
"""create dxb_kernel"""
|
||||
def __init__(self):
|
||||
super(dxb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((1, 0), (0, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
x = x - x1
|
||||
return x
|
||||
|
||||
|
||||
class dxf_kernel(nn.Cell):
|
||||
"""create dxf_kernel"""
|
||||
def __init__(self):
|
||||
super(dxf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 1), (0, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (1, 0, 0), x_shape)
|
||||
x = x1 - x
|
||||
return x
|
||||
|
||||
|
||||
class dyb_kernel(nn.Cell):
|
||||
"""create dyb_kernel"""
|
||||
def __init__(self):
|
||||
super(dyb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (1, 0), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
x = x - x1
|
||||
return x
|
||||
|
||||
|
||||
class dyf_kernel(nn.Cell):
|
||||
"""create dyf_kernel"""
|
||||
def __init__(self):
|
||||
super(dyf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 1), (0, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 1, 0), x_shape)
|
||||
x = x1 - x
|
||||
return x
|
||||
|
||||
|
||||
class dzb_kernel(nn.Cell):
|
||||
"""create dzb_kernel"""
|
||||
def __init__(self):
|
||||
super(dzb_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 0), (1, 0)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 0), x_shape)
|
||||
x = x - x1
|
||||
return x
|
||||
|
||||
|
||||
class dzf_kernel(nn.Cell):
|
||||
"""create dzf_kernel"""
|
||||
def __init__(self):
|
||||
super(dzf_kernel, self).__init__()
|
||||
self.pad = P.Pad(((0, 0), (0, 0), (0, 1)))
|
||||
self.slice = P.Slice()
|
||||
self.shape = P.Shape()
|
||||
|
||||
def construct(self, x):
|
||||
x1 = self.pad(x)
|
||||
x_shape = self.shape(x)
|
||||
x1 = self.slice(x1, (0, 0, 1), x_shape)
|
||||
x = x1 - x
|
||||
return x
|
||||
@ -0,0 +1,43 @@
|
||||
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# 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.
|
||||
# ============================================================================
|
||||
"""read variables"""
|
||||
|
||||
import numpy as np
|
||||
import netCDF4 as nc
|
||||
|
||||
|
||||
# variable name list
|
||||
params_name = ['z', 'zz', 'dz', 'dzz', 'dx', 'dy', 'cor', 'h', 'fsm', 'dum', 'dvm', 'art', 'aru', 'arv', 'rfe', 'rfw',
|
||||
'rfn', 'rfs', 'east_e', 'north_e', 'east_c', 'north_c', 'east_u', 'north_u', 'east_v', 'north_v', 'tb',
|
||||
'sb', 'tclim', 'sclim', 'rot', 'vfluxf', 'wusurf', 'wvsurf', 'e_atmos', 'ub', 'vb', 'uab', 'vab', 'elb',
|
||||
'etb', 'dt', 'uabw', 'uabe', 'vabs', 'vabn', 'els', 'eln', 'ele', 'elw', 'ssurf', 'tsurf', 'tbe', 'sbe',
|
||||
'sbw', 'tbw', 'tbn', 'tbs', 'sbn', 'sbs', 'wtsurf', 'swrad']
|
||||
|
||||
|
||||
def load_var(file_obj, name):
|
||||
"""load variable from nc data file"""
|
||||
data = file_obj.variables[name]
|
||||
data = data[:]
|
||||
data = np.float32(np.transpose(data, (2, 1, 0)))
|
||||
return data
|
||||
|
||||
|
||||
def read_nc(file_path):
|
||||
""" put the load variable into the dict """
|
||||
variable = {}
|
||||
file_obj = nc.Dataset(file_path)
|
||||
for name in params_name:
|
||||
variable[name] = load_var(file_obj, name)
|
||||
return variable
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,71 @@
|
||||
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# 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.
|
||||
# ============================================================================
|
||||
"""train"""
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
import mindspore.context as context
|
||||
from src.read_var import read_nc
|
||||
from src.GOMO import GOMO_init, GOMO, read_init
|
||||
|
||||
parser = argparse.ArgumentParser(description='GOMO')
|
||||
parser.add_argument('--file_path', type=str, default=None, help='file path')
|
||||
parser.add_argument('--outputs_path', type=str, default=None, help='outputs path')
|
||||
parser.add_argument('--im', type=int, default=65, help='im size')
|
||||
parser.add_argument('--jm', type=int, default=49, help='jm size')
|
||||
parser.add_argument('--kb', type=int, default=21, help='kb size')
|
||||
parser.add_argument('--stencil_width', type=int, default=1, help='stencil width')
|
||||
parser.add_argument('--step', type=int, default=10, help='time step')
|
||||
args_gomo = parser.parse_args()
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="GPU", save_graphs=False, enable_graph_kernel=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
variable = read_nc(args_gomo.file_path)
|
||||
im = args_gomo.im
|
||||
jm = args_gomo.jm
|
||||
kb = args_gomo.kb
|
||||
stencil_width = args_gomo.stencil_width
|
||||
|
||||
# variable init
|
||||
dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, dt, h, w, wubot, wvbot, vfluxb, utb, vtb, dhb, egb, vfluxf, z, zz, \
|
||||
dzz, cor, fsm = read_init(
|
||||
variable, im, jm, kb)
|
||||
|
||||
# define grid and init variable update
|
||||
net_init = GOMO_init(im, jm, kb, stencil_width)
|
||||
ua, va, el, et, etf, d, dt, l, q2b, q2lb, kh, km, kq, aam, w, q2, q2l, t, s, u, v, cbc, rmean, rho, x_d, y_d, z_d\
|
||||
= net_init(dx, dy, dz, uab, vab, elb, etb, sb, tb, ub, vb, h, w, vfluxf, zz, fsm)
|
||||
|
||||
# define GOMO model
|
||||
Model = GOMO(im=im, jm=jm, kb=kb, stencil_width=stencil_width, variable=variable, x_d=x_d, y_d=y_d, z_d=z_d,
|
||||
q2b=q2b, q2lb=q2lb, aam=aam, cbc=cbc, rmean=rmean)
|
||||
|
||||
# time step of GOMO Model
|
||||
for step in range(1, args_gomo.step+1):
|
||||
elf, etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, ub, vb, \
|
||||
egb, etb, dt, dhb, utb, vtb, vfluxb, et, steps, vamax, q2b, q2lb = Model(
|
||||
etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho,
|
||||
wubot, wvbot, ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et)
|
||||
vars_list = etf, ua, uab, va, vab, el, elb, d, u, v, w, kq, km, kh, q2, q2l, tb, t, sb, s, rho, wubot, wvbot, \
|
||||
ub, vb, egb, etb, dt, dhb, utb, vtb, vfluxb, et
|
||||
for var in vars_list:
|
||||
var.asnumpy()
|
||||
# save output
|
||||
if step % 5 == 0:
|
||||
np.save(args_gomo.outputs_path + "u_"+str(step)+".npy", u.asnumpy())
|
||||
np.save(args_gomo.outputs_path + "v_" + str(step) + ".npy", v.asnumpy())
|
||||
np.save(args_gomo.outputs_path + "t_" + str(step) + ".npy", t.asnumpy())
|
||||
np.save(args_gomo.outputs_path + "et_" + str(step) + ".npy", et.asnumpy())
|
||||
Loading…
Reference in new issue