In fluid simulation problems, capturing the complex details of turbulence has always been a long-standing challenge for numerical simulation. Solving these details with discrete models will incur huge computational costs, which will soon become infeasible for flows on human spatial and temporal scales. Therefore, the demand for fluid super-resolution has emerged, which aims to recover high-resolution fluid simulation results from low-resolution results through fluid dynamics simulation and deep learning technology, so as to reduce the huge computational cost in the process of generating high-resolution fluids. This technology can be applied to various fluid simulations, such as water flow, air flow, flame simulation, etc.
Generative Adversarial Networks (GAN) is a deep learning network using unsupervised learning methods. GAN networks (at least) contain two models: Generator and Discriminator. The generator is used to generate the output of the problem, and the discriminator is used to judge whether the output is true or false. Both optimize together in mutual game, and finally make the output of the generator close to the true value.
Based on the GAN network, tempoGAN adds a time-related discriminator Discriminator_tempo. The network structure of this discriminator is the same as the basic discriminator, but the input is several consecutive frames of data in time, rather than a single frame of data, thereby taking timing into consideration.
This problem mainly uses this network to obtain corresponding high-density fluid data through input low-density fluid data, greatly saving time and computational costs.
This problem includes three models: Generator, Discriminator and time-related discriminator (Discriminator_tempo). According to the training process of the GAN network, these three models are trained alternately, and the training order is: Discriminator, Discriminator_tempo, Generator.
GAN network is unsupervised learning. In the network design of this problem, the target value is used as an input value and input into the network for training.
Next, we will explain how to convert the problem into PaddleScience code step by step and solve the problem using deep learning methods. In order to quickly understand PaddleScience, only key steps such as model construction and constraint construction are described below, while other details please refer to API Documentation.
The dataset is a 2d fluid dataset generated using the open source code package mantaflow. The dataset includes numerical values converted from low and high-density fluid images of a certain number of continuous frames, stored in dictionary form in .mat files.
Before running the code for this problem, please download training dataset and validation dataset, and store them separately in the path after downloading:
The figure above is the complete model structure diagram of tempoGAN, but this problem only deals with relatively simple cases, and does not involve parts including velocity and vorticity input, 3d, data augmentation, advection operator, etc. If you are interested in contents not included in these documents, you can modify the code yourself and conduct further experiments.
As shown in the figure above, the input of Generator is the interpolation of low-density fluid data, and the output is the generated high-density fluid simulation data. The input of Discriminator is the concatenation of interpolation of low-density fluid data with high-density fluid simulation data generated by Generator and target high-density fluid data respectively. The input of Discriminator_tempo is multi-frame continuous high-density fluid simulation data generated by Generator and target high-density fluid data.
Although the composition of input and output looks complicated, they are essentially fluid density data, so the mapping functions of the 3 networks are all \(f: \mathbb{R}^1 \to \mathbb{R}^1\).
Different from simple MLP networks, depending on different problems to be solved, GAN generators and discriminators have a variety of network structures to choose from, which will not be repeated here. Due to this uniqueness, the tempoGAN network in this problem is not built into PaddleScience and needs to be implemented additionally.
The Generator in this problem is a model with 4 layers of improved Res Block, Discriminator and Discriminator_tempo are the same model with 4 layers of convolution results, both have the same network structure but different inputs. The network parameters of Generator, Discriminator and Discriminator_tempo also need to be defined additionally.
For specific code, please refer to the gan.py file in Complete Code.
Since the intermediate results of the generator and discriminator in the GAN network need to be called mutually and participate in each other's loss calculation, Model List is used for implementation, expressed in PaddleScience code as follows:
Note that the network input defined in the above code is not exactly the same as the actual network input, so transform needs to be performed on the input.
The input of Generator is the interpolation of low-density fluid data, while the dataset stores the original low-density fluid data, so an interpolation transform is required.
Indicates stopping the calculation gradient of parameters. This is set because this variable is only used as input for Discriminator and Discriminator_tempo here, and should not participate in gradient backpropagation during reverse calculation. If such setting is not made, since this variable comes from the output of Generator, the gradient will be transmitted to Generator along this variable during backpropagation, thereby changing the parameters in Generator, which is obviously not what we want.
In this way, we instantiate a neural network model model list possessing Generator, Discriminator and Discriminator_tempo and containing input transform.
log_freq:20DATASET_PATH:./datasets/tempoGAN/2d_train.matDATASET_PATH_VALID:./datasets/tempoGAN/2d_valid.mat# set working conditionUSE_AMP:trueUSE_SPATIALDISC:trueUSE_TEMPODISC:trueWEIGHT_GEN:[5.0,0.0,1.0]# lambda_l1, lambda_l2, lambda_tWEIGHT_GEN_LAYER:[-1.0e-5,-1.0e-5,-1.0e-5,-1.0e-5,-1.0e-5]WEIGHT_DISC:1.0
Note that it contains 3 bool type variables use_amp, use_spatialdisc and use_tempodisc, which respectively represent whether to use mixed precision training (AMP), whether to use Discriminator and whether to use Discriminator_tempo. When both use_spatialdisc and use_tempodisc are set to False, the network structure of this problem will become a pure Generator model and is no longer a GAN network.
At the same time, hyperparameters such as training epochs and learning rate need to be specified. Note that since the GAN network training process is different from general single-model networks, the setting of EPOCHS is also different.
The training uses the Adam optimizer, and the learning rate is reduced to \(1/20\) of the original when Epoch reaches half, so the Step method is used as the learning rate strategy. If by_epoch is set to True, the learning rate will change according to the training Epoch, otherwise it will change according to Iteration.
# initialize Adam optimizerlr_scheduler_gen=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_gen=ppsci.optimizer.Adam(lr_scheduler_gen)(model_gen)ifcfg.USE_SPATIALDISC:lr_scheduler_disc=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_disc=ppsci.optimizer.Adam(lr_scheduler_disc)(model_disc)ifcfg.USE_TEMPODISC:lr_scheduler_disc_tempo=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_disc_tempo=ppsci.optimizer.Adam(lr_scheduler_disc_tempo)((model_disc_tempo,))
This problem adopts unsupervised learning method. Although it is not trained in a supervised learning manner, supervised constraint SupervisedConstraint can still be used here. Before defining constraints, data reading configurations such as file path need to be specified for supervised constraints. Since tempoGAN belongs to self-supervised learning, there is no label data in the dataset, but a part of input data is used as label, so output_expr of the constraint needs to be set.
The first parameter of SupervisedConstraint is the reading configuration of supervised constraint, where dataset field represents the training dataset information used, and each field respectively represents:
name: Dataset type, here NamedArrayDataset represents dataset of .mat type read from Array;
input: Input data of Array type;
label: Label data of Array type;
transforms: All data transform methods, here FunctionalTransform is a custom data transform class reserved by PaddleScience, which supports custom transform of input data when writing code. For specific code, please refer to Custom loss and data transform;
batch_size field represents the size of batch;
sampler field represents sampling method, where each field represents:
name: Sampler type, here BatchSampler represents batch sampler;
drop_last: Whether to discard the last samples that cannot make up a mini-batch, default is False;
shuffle: Whether to shuffle the order when generating sample subscripts, default is False;
The second parameter is the loss function. Here FunctionalLoss is a custom loss function class reserved by PaddleScience, which supports custom loss calculation method when writing code, rather than using existing methods such as MSE. For specific code, please refer to Custom loss and data transform.
The third parameter is output_expr of the constraint condition. As mentioned above, it is to allow the program to use input data as label.
The fourth parameter is the name of the constraint condition. We need to name each constraint condition for subsequent indexing.
After the constraints are constructed, encapsulate them into a dictionary with the name we just named as the keyword for subsequent access. Since use_spatialdisc and use_tempodisc are set in this problem, some constraints of Generator may not exist, so first encapsulate the certainly existing constraints into the dictionary, and when other constraints exist, add constraint elements to the dictionary.
Because of the characteristics of GAN network training, this problem does not use the built-in visualizer in PaddleScience, but customizes a function for implementing inference. This function reads validation set data, obtains inference results and saves the results in image form. Calling this function at regular intervals during the training process can monitor the training effect during the training process.
defpredict_and_save_plot(output_dir:str,epoch_id:int,solver_gen:ppsci.solver.Solver,dataset_valid:np.ndarray,tile_ratio:int=1,):"""Predicting and plotting. Args: output_dir (str): Output dir path. epoch_id (int): Which epoch it is. solver_gen (ppsci.solver.Solver): Solver for predicting. dataset_valid (np.ndarray): Valid dataset. tile_ratio (int, optional): How many tiles of one dim. Defaults to 1. """dir_pred="predict/"os.makedirs(os.path.join(output_dir,dir_pred),exist_ok=True)start_idx=190density_low=dataset_valid["density_low"][start_idx:start_idx+3]density_high=dataset_valid["density_high"][start_idx:start_idx+3]# tiledensity_low=(split_data(density_low,tile_ratio)iftile_ratio!=1elsedensity_low)density_high=(split_data(density_high,tile_ratio)iftile_ratio!=1elsedensity_high)pred_dict=solver_gen.predict({"density_low":density_low,"density_high":density_high,},{"density_high":lambdaout:out["output_gen"]},batch_size=tile_ratio*tile_ratioiftile_ratio!=1else3,no_grad=False,)ifepoch_id==1:# plot interpolated input imageinput_img=np.expand_dims(dataset_valid["density_low"][start_idx],axis=0)input_img=paddle.to_tensor(input_img,dtype=paddle.get_default_dtype())input_img=F.interpolate(input_img,[input_img.shape[-2]*4,input_img.shape[-1]*4],mode="nearest",).numpy()Img.imsave(os.path.join(output_dir,dir_pred,"input.png"),np.squeeze(input_img),vmin=0.0,vmax=1.0,cmap="gray",)# plot target imageImg.imsave(os.path.join(output_dir,dir_pred,"target.png"),np.squeeze(dataset_valid["density_high"][start_idx]),vmin=0.0,vmax=1.0,cmap="gray",)# plot pred imagepred_img=(concat_data(pred_dict["density_high"].numpy(),tile_ratio)iftile_ratio!=1elsenp.squeeze(pred_dict["density_high"][0].numpy()))Img.imsave(os.path.join(output_dir,dir_pred,f"pred_epoch_{str(epoch_id)}.png"),pred_img,vmin=0.0,vmax=1.0,cmap="gray",)
Since this problem adopts unsupervised learning and there is no label data in the data, loss is calculated, so loss needs to be customized. The method is to define relevant functions first, and then pass the function name as a parameter to FunctionalLoss. It should be noted that the input and output parameters of the custom loss function need to be consistent with other functions such as MSE in PaddleScience, that is, the input is dictionary variables such as model output output_dict, and the output is loss value paddle.Tensor.
The loss of Generator provides l1 loss, l2 loss, loss judged by Discriminator on output and loss judged by Discriminator_tempo on output. Whether these losses exist is controlled according to weight parameters. If the weight parameter of a certain loss item is 0, it means that the loss item is not added during training.
defloss_func_gen(self,output_dict:Dict,*args)->paddle.Tensor:"""Calculate loss of generator when use spatial discriminator. The loss consists of l1 loss, l2 loss and layer loss when use spatial discriminator. Notice that all item of loss is optional because weight of them might be 0. Args: output_dict (Dict): output dict of model. Returns: paddle.Tensor: Loss of generator. """# l1 lossloss_l1=F.l1_loss(output_dict["output_gen"],output_dict["density_high"],"mean")losses=loss_l1*self.weight_gen[0]# l2 lossloss_l2=F.mse_loss(output_dict["output_gen"],output_dict["density_high"],"mean")losses+=loss_l2*self.weight_gen[1]ifself.weight_gen_layerisnotNone:# disc(generator_out) lossout_disc_from_gen=output_dict["out_disc_from_gen"][-1]label_ones=paddle.ones_like(out_disc_from_gen)loss_gen=F.binary_cross_entropy_with_logits(out_disc_from_gen,label_ones,reduction="mean")losses+=loss_gen# layer losskey_list=list(output_dict.keys())# ["out0_layer0","out0_layer1","out0_layer2","out0_layer3","out_disc_from_target",# "out1_layer0","out1_layer1","out1_layer2","out1_layer3","out_disc_from_gen"]loss_layer=0foriinrange(1,len(self.weight_gen_layer)):# i = 0,1,2,3loss_layer+=(self.weight_gen_layer[i]*F.mse_loss(output_dict[key_list[i]],output_dict[key_list[5+i]],reduction="sum",)/2)losses+=loss_layer*self.weight_gen_layer[0]return{"output_gen":losses}defloss_func_gen_tempo(self,output_dict:Dict,*args)->paddle.Tensor:"""Calculate loss of generator when use temporal discriminator. The loss is cross entropy loss when use temporal discriminator. Args: output_dict (Dict): output dict of model. Returns: paddle.Tensor: Loss of generator. """out_disc_tempo_from_gen=output_dict["out_disc_tempo_from_gen"][-1]label_t_ones=paddle.ones_like(out_disc_tempo_from_gen)loss_gen_t=F.binary_cross_entropy_with_logits(out_disc_tempo_from_gen,label_t_ones,reduction="mean")losses=loss_gen_t*self.weight_gen[2]return{"out_disc_tempo_from_gen":losses}
Discriminator is a discriminator, its function is to judge whether data is true data or false data, so its loss is the loss generated by judging data generated by Generator as false and the loss generated by judging target value data as true.
This problem provides an input data processing method, randomly cropping a piece of input fluid density data, then judging the density value. If the density value of the cropped block is lower than the threshold, it is re-cropped until the density meets the condition or the number of cropping times reaches the threshold. This is mainly done to reduce the video memory required for training, and the judgment of the block density value ensures the richness of information in the block. In Parameter and Hyperparameter Setting, tile_ratio indicates how many times the original size is compared to the block size, that is, if tile_ratio is 2, the size of the cropped block is one-fourth of the entire original image.
classDataFuncs:"""All functions used for data transform. Args: tile_ratio (int, optional): How many tiles of one dim. Defaults to 1. density_min (float, optional): Minimize density of one tile. Defaults to 0.02. max_turn (int, optional): Maximize turn of taking a tile from one image. Defaults to 20. """def__init__(self,tile_ratio:int=1,density_min:float=0.02,max_turn:int=20)->None:self.tile_ratio=tile_ratioself.density_min=density_minself.max_turn=max_turndeftransform(self,input_item:Dict[str,np.ndarray],label_item:Dict[str,np.ndarray],weight_item:Dict[str,np.ndarray],)->Tuple[Dict[str,np.ndarray],Dict[str,np.ndarray],Dict[str,np.ndarray]]:ifself.tile_ratio==1:returninput_item,label_item,weight_itemfor_inrange(self.max_turn):rand_ratio=np.random.rand()density_low=self.cut_data(input_item["density_low"],rand_ratio)density_high=self.cut_data(input_item["density_high"],rand_ratio)ifself.is_valid_tile(density_low):breakinput_item["density_low"]=density_lowinput_item["density_high"]=density_highreturninput_item,label_item,weight_itemdefcut_data(self,data:np.ndarray,rand_ratio:float)->paddle.Tensor:# data: C,H,W_,H,W=data.shapeifH%self.tile_ratio!=0orW%self.tile_ratio!=0:exit(f"ERROR: input images cannot be divided into {self.tile_ratio} parts evenly!")tile_shape=[H//self.tile_ratio,W//self.tile_ratio]rand_shape=np.floor(rand_ratio*(np.array([H,W])-np.array(tile_shape)))start=[int(rand_shape[0]),int(rand_shape[1])]end=[int(rand_shape[0]+tile_shape[0]),int(rand_shape[1]+tile_shape[1])]data=paddle.slice(paddle.to_tensor(data),axes=[-2,-1],starts=start,ends=end)returndatadefis_valid_tile(self,tile:paddle.Tensor):img_density=tile[0].sum()returnimg_density>=(self.density_min*tile.shape[0]*tile.shape[1]*tile.shape[2])
Note that the code here only provides the idea of data transform. The simple block method in the current code will obviously affect the training effect due to less information contained in the input. Therefore, in this problem, when the video memory is sufficient, tile_ratio should be set to 1. When the video memory is insufficient, it is also recommended to prioritize using mixed precision training to reduce current memory usage.
Note that the GAN type network training method is alternating training of multiple models, which is different from single model or multi-model phased training, and solver.train API cannot be simply used. For specific code, please refer to tempoGAN.py file in Complete Code.
During training, only target results and model output results of specific images are saved at specific Epoch, and an evaluation is performed on the output result of the last Epoch after training ends, so as to intuitively evaluate the model optimization effect. Do not use the built-in evaluator in PaddleScience, nor evaluate during the training process:
foriinrange(1,cfg.TRAIN.epochs+1):logger.message(f"\nEpoch: {i}\n")# plotting during trainingifi==1ori%PRED_INTERVAL==0ori==cfg.TRAIN.epochs:func_module.predict_and_save_plot(cfg.output_dir,i,solver_gen,dataset_valid,cfg.TILE_RATIO)
############### evaluation for training ###############img_target=(func_module.get_image_array(os.path.join(cfg.output_dir,"predict","target.png"))/255.0)img_pred=(func_module.get_image_array(os.path.join(cfg.output_dir,"predict",f"pred_epoch_{cfg.TRAIN.epochs}.png"))/255.0)eval_mse,eval_psnr,eval_ssim=func_module.evaluate_img(img_target,img_pred)logger.message(f"MSE: {eval_mse}, PSNR: {eval_psnr}, SSIM: {eval_ssim}")
For specific code, please refer to tempoGAN.py file in Complete Code.
The evaluation metric for this problem is to compare the super-resolution result output by the model with the actual high-resolution image, and use three indicators MSE (Mean-Square Error), PSNR (Peak Signal-to-Noise Ratio), and SSIM (Structural SIMilarity) to evaluate image similarity. Therefore, the built-in evaluator in PaddleScience is not used, nor is there a Solver.eval() process.
Provides the option to save the model output result, so as to see the result after super-resolution more intuitively. Whether to open is specified by save_outs in configuration file EVAL:
The complete code contains PaddleScience specific training process code tempoGAN.py and all custom function code functions.py. In addition, network structure code gan.py is added to ppsci.arch, which is shown below together. If you need to customize the network structure, you can use it as a reference.
# Copyright (c) 2023 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.importosfromosimportpathasospimportfunctionsasfunc_moduleimporthydraimportnumpyasnpimportpaddlefromomegaconfimportDictConfigimportppscifromppsci.utilsimportcheckerfromppsci.utilsimportloggerfromppsci.utilsimportsave_loadifnotchecker.dynamic_import_to_globals("hdf5storage"):raiseImportError("Could not import hdf5storage python package. ""Please install it with `pip install hdf5storage`.")importhdf5storagedeftrain(cfg:DictConfig):ppsci.utils.misc.set_random_seed(cfg.seed)# initialize loggerlogger.init_logger("ppsci",osp.join(cfg.output_dir,"train.log"),"info")gen_funcs=func_module.GenFuncs(cfg.WEIGHT_GEN,(cfg.WEIGHT_GEN_LAYERifcfg.USE_SPATIALDISCelseNone))disc_funcs=func_module.DiscFuncs(cfg.WEIGHT_DISC)data_funcs=func_module.DataFuncs(cfg.TILE_RATIO)# load datasetlogger.message("Attention! Start loading datasets, this will take tens of seconds to several minutes, please wait patiently.")dataset_train=hdf5storage.loadmat(cfg.DATASET_PATH)logger.message("Finish loading training dataset.")dataset_valid=hdf5storage.loadmat(cfg.DATASET_PATH_VALID)logger.message("Finish loading validation dataset.")# define Generator modelmodel_gen=ppsci.arch.Generator(**cfg.MODEL.gen_net)model_gen.register_input_transform(gen_funcs.transform_in)disc_funcs.model_gen=model_genmodel_tuple=(model_gen,)# define Discriminatorsifcfg.USE_SPATIALDISC:model_disc=ppsci.arch.Discriminator(**cfg.MODEL.disc_net)model_disc.register_input_transform(disc_funcs.transform_in)model_tuple+=(model_disc,)# define temporal Discriminatorsifcfg.USE_TEMPODISC:model_disc_tempo=ppsci.arch.Discriminator(**cfg.MODEL.tempo_net)model_disc_tempo.register_input_transform(disc_funcs.transform_in_tempo)model_tuple+=(model_disc_tempo,)# define model_listmodel_list=ppsci.arch.ModelList(model_tuple)# initialize Adam optimizerlr_scheduler_gen=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_gen=ppsci.optimizer.Adam(lr_scheduler_gen)(model_gen)ifcfg.USE_SPATIALDISC:lr_scheduler_disc=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_disc=ppsci.optimizer.Adam(lr_scheduler_disc)(model_disc)ifcfg.USE_TEMPODISC:lr_scheduler_disc_tempo=ppsci.optimizer.lr_scheduler.Step(step_size=cfg.TRAIN.epochs//2,**cfg.TRAIN.lr_scheduler)()optimizer_disc_tempo=ppsci.optimizer.Adam(lr_scheduler_disc_tempo)((model_disc_tempo,))# Generator# manually build constraint(s)sup_constraint_gen=ppsci.constraint.SupervisedConstraint({"dataset":{"name":"NamedArrayDataset","input":{"density_low":dataset_train["density_low"],"density_high":dataset_train["density_high"],},"transforms":({"FunctionalTransform":{"transform_func":data_funcs.transform,},},),},"batch_size":cfg.TRAIN.batch_size.sup_constraint,"sampler":{"name":"BatchSampler","drop_last":False,"shuffle":False,},},ppsci.loss.FunctionalLoss(gen_funcs.loss_func_gen),{"output_gen":lambdaout:out["output_gen"],"density_high":lambdaout:out["density_high"],},name="sup_constraint_gen",)constraint_gen={sup_constraint_gen.name:sup_constraint_gen}ifcfg.USE_TEMPODISC:sup_constraint_gen_tempo=ppsci.constraint.SupervisedConstraint({"dataset":{"name":"NamedArrayDataset","input":{"density_low":dataset_train["density_low_tempo"],"density_high":dataset_train["density_high_tempo"],},"transforms":({"FunctionalTransform":{"transform_func":data_funcs.transform,},},),},"batch_size":int(cfg.TRAIN.batch_size.sup_constraint//3),"sampler":{"name":"BatchSampler","drop_last":False,"shuffle":False,},},ppsci.loss.FunctionalLoss(gen_funcs.loss_func_gen_tempo),{"output_gen":lambdaout:out["output_gen"],"density_high":lambdaout:out["density_high"],},name="sup_constraint_gen_tempo",)constraint_gen[sup_constraint_gen_tempo.name]=sup_constraint_gen_tempo# Discriminators# manually build constraint(s)ifcfg.USE_SPATIALDISC:sup_constraint_disc=ppsci.constraint.SupervisedConstraint({"dataset":{"name":"NamedArrayDataset","input":{"density_low":dataset_train["density_low"],"density_high":dataset_train["density_high"],},"label":{"out_disc_from_target":np.ones((np.shape(dataset_train["density_high"])[0],1),dtype=paddle.get_default_dtype(),),"out_disc_from_gen":np.ones((np.shape(dataset_train["density_high"])[0],1),dtype=paddle.get_default_dtype(),),},"transforms":({"FunctionalTransform":{"transform_func":data_funcs.transform,},},),},"batch_size":cfg.TRAIN.batch_size.sup_constraint,"sampler":{"name":"BatchSampler","drop_last":False,"shuffle":False,},},ppsci.loss.FunctionalLoss(disc_funcs.loss_func),name="sup_constraint_disc",)constraint_disc={sup_constraint_disc.name:sup_constraint_disc}# temporal Discriminators# manually build constraint(s)ifcfg.USE_TEMPODISC:sup_constraint_disc_tempo=ppsci.constraint.SupervisedConstraint({"dataset":{"name":"NamedArrayDataset","input":{"density_low":dataset_train["density_low_tempo"],"density_high":dataset_train["density_high_tempo"],},"label":{"out_disc_tempo_from_target":np.ones((np.shape(dataset_train["density_high_tempo"])[0],1),dtype=paddle.get_default_dtype(),),"out_disc_tempo_from_gen":np.ones((np.shape(dataset_train["density_high_tempo"])[0],1),dtype=paddle.get_default_dtype(),),},"transforms":({"FunctionalTransform":{"transform_func":data_funcs.transform,},},),},"batch_size":int(cfg.TRAIN.batch_size.sup_constraint//3),"sampler":{"name":"BatchSampler","drop_last":False,"shuffle":False,},},ppsci.loss.FunctionalLoss(disc_funcs.loss_func_tempo),name="sup_constraint_disc_tempo",)constraint_disc_tempo={sup_constraint_disc_tempo.name:sup_constraint_disc_tempo}# initialize solversolver_gen=ppsci.solver.Solver(model_list,constraint_gen,cfg.output_dir,optimizer_gen,lr_scheduler_gen,cfg.TRAIN.epochs_gen,cfg.TRAIN.iters_per_epoch,eval_during_train=cfg.TRAIN.eval_during_train,use_amp=cfg.USE_AMP,amp_level=cfg.TRAIN.amp_level,)ifcfg.USE_SPATIALDISC:solver_disc=ppsci.solver.Solver(model_list,constraint_disc,cfg.output_dir,optimizer_disc,lr_scheduler_disc,cfg.TRAIN.epochs_disc,cfg.TRAIN.iters_per_epoch,eval_during_train=cfg.TRAIN.eval_during_train,use_amp=cfg.USE_AMP,amp_level=cfg.TRAIN.amp_level,)ifcfg.USE_TEMPODISC:solver_disc_tempo=ppsci.solver.Solver(model_list,constraint_disc_tempo,cfg.output_dir,optimizer_disc_tempo,lr_scheduler_disc_tempo,cfg.TRAIN.epochs_disc_tempo,cfg.TRAIN.iters_per_epoch,eval_during_train=cfg.TRAIN.eval_during_train,use_amp=cfg.USE_AMP,amp_level=cfg.TRAIN.amp_level,)PRED_INTERVAL=200foriinrange(1,cfg.TRAIN.epochs+1):logger.message(f"\nEpoch: {i}\n")# plotting during trainingifi==1ori%PRED_INTERVAL==0ori==cfg.TRAIN.epochs:func_module.predict_and_save_plot(cfg.output_dir,i,solver_gen,dataset_valid,cfg.TILE_RATIO)disc_funcs.model_gen=model_gen# train disc, input: (x,y,G(x))ifcfg.USE_SPATIALDISC:solver_disc.train()# train disc tempo, input: (y_3,G(x)_3)ifcfg.USE_TEMPODISC:solver_disc_tempo.train()# train gen, input: (x,)solver_gen.train()############### evaluation for training ###############img_target=(func_module.get_image_array(os.path.join(cfg.output_dir,"predict","target.png"))/255.0)img_pred=(func_module.get_image_array(os.path.join(cfg.output_dir,"predict",f"pred_epoch_{cfg.TRAIN.epochs}.png"))/255.0)eval_mse,eval_psnr,eval_ssim=func_module.evaluate_img(img_target,img_pred)logger.message(f"MSE: {eval_mse}, PSNR: {eval_psnr}, SSIM: {eval_ssim}")defevaluate(cfg:DictConfig):ifcfg.EVAL.save_outs:frommatplotlibimportimageasImgos.makedirs(osp.join(cfg.output_dir,"eval_outs"),exist_ok=True)ppsci.utils.misc.set_random_seed(cfg.seed)# initialize loggerlogger.init_logger("ppsci",osp.join(cfg.output_dir,"eval.log"),"info")gen_funcs=func_module.GenFuncs(cfg.WEIGHT_GEN,None)# load datasetdataset_valid=hdf5storage.loadmat(cfg.DATASET_PATH_VALID)# define Generator modelmodel_gen=ppsci.arch.Generator(**cfg.MODEL.gen_net)model_gen.register_input_transform(gen_funcs.transform_in)# define model_listmodel_list=ppsci.arch.ModelList((model_gen,))# load pretrained modelsave_load.load_pretrain(model_list,cfg.EVAL.pretrained_model_path)# set validatoreval_dataloader_cfg={"dataset":{"name":"NamedArrayDataset","input":{"density_low":dataset_valid["density_low"],},"label":{"density_high":dataset_valid["density_high"]},},"sampler":{"name":"BatchSampler","drop_last":False,"shuffle":False,},"batch_size":1,}sup_validator=ppsci.validate.SupervisedValidator(eval_dataloader_cfg,ppsci.loss.MSELoss("mean"),{"density_high":lambdaout:out["output_gen"]},metric={"metric":ppsci.metric.L2Rel()},name="sup_validator_gen",)# customized evaluationdefscale(data):smax=np.max(data)smin=np.min(data)return(data-smin)/(smax-smin)eval_mse_list=[]eval_psnr_list=[]eval_ssim_list=[]fori,(input,label,_)inenumerate(sup_validator.data_loader):output_dict=model_list({"density_low":input["density_low"]})output_arr=scale(np.squeeze(output_dict["output_gen"].numpy()))target_arr=scale(np.squeeze(label["density_high"].numpy()))eval_mse,eval_psnr,eval_ssim=func_module.evaluate_img(target_arr,output_arr)eval_mse_list.append(eval_mse)eval_psnr_list.append(eval_psnr)eval_ssim_list.append(eval_ssim)ifcfg.EVAL.save_outs:Img.imsave(osp.join(cfg.output_dir,"eval_outs",f"out_{i}.png"),output_arr,vmin=0.0,vmax=1.0,cmap="gray",)logger.message(f"MSE: {np.mean(eval_mse_list)}, PSNR: {np.mean(eval_psnr_list)}, SSIM: {np.mean(eval_ssim_list)}")defexport(cfg:DictConfig):frompaddle.staticimportInputSpec# set modelsgen_funcs=func_module.GenFuncs(cfg.WEIGHT_GEN,None)model_gen=ppsci.arch.Generator(**cfg.MODEL.gen_net)model_gen.register_input_transform(gen_funcs.transform_in)# define model_listmodel_list=ppsci.arch.ModelList((model_gen,))# load pretrained modelsolver=ppsci.solver.Solver(model=model_list,pretrained_model_path=cfg.INFER.pretrained_model_path)# export modelsinput_spec=[{"density_low":InputSpec([None,1,128,128],"float32",name="density_low")},]solver.export(input_spec,cfg.INFER.export_path,skip_prune_program=True)definference(cfg:DictConfig):frommatplotlibimportimageasImgfromdeploy.python_inferimportpinn_predictor# set model predictorpredictor=pinn_predictor.PINNPredictor(cfg)# load datasetdataset_infer={"density_low":hdf5storage.loadmat(cfg.DATASET_PATH_VALID)["density_low"]}output_dict=predictor.predict(dataset_infer,cfg.INFER.batch_size)# mapping data to cfg.INFER.output_keysoutput=[output_dict[key]forkeyinoutput_dict]defscale(data):smax=np.max(data)smin=np.min(data)return(data-smin)/(smax-smin)fori,imginenumerate(output[0]):img=scale(np.squeeze(img))Img.imsave(osp.join(cfg.output_dir,f"out_{i}.png"),img,vmin=0.0,vmax=1.0,cmap="gray",)@hydra.main(version_base=None,config_path="./conf",config_name="tempogan.yaml")defmain(cfg:DictConfig):ifcfg.mode=="train":train(cfg)elifcfg.mode=="eval":evaluate(cfg)elifcfg.mode=="export":export(cfg)elifcfg.mode=="infer":inference(cfg)else:raiseValueError(f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'")if__name__=="__main__":main()
# Copyright (c) 2023 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.importosfromtypingimportDictfromtypingimportListfromtypingimportTupleimportnumpyasnpimportpaddleimportpaddle.nn.functionalasFfrommatplotlibimportimageasImgfromPILimportImagefromskimage.metricsimportmean_squared_errorfromskimage.metricsimportpeak_signal_noise_ratiofromskimage.metricsimportstructural_similarityimportppscifromppsci.utilsimportlogger# traindefinterpolate(data:paddle.Tensor,ratio:int,mode:str="nearest")->paddle.Tensor:"""Interpolate twice. Args: data (paddle.Tensor): The data to be interpolated. ratio (int): Ratio of one interpolation. mode (str, optional): Interpolation method. Defaults to "nearest". Returns: paddle.Tensor: Data interpolated. """for_inrange(2):data=F.interpolate(data,[data.shape[-2]*ratio,data.shape[-1]*ratio],mode=mode,)returndatadefreshape_input(input_dict:Dict[str,paddle.Tensor])->Dict[str,paddle.Tensor]:"""Reshape input data for temporally Discriminator. Reshape data from N, C, W, H to N * C, 1, H, W. Which will merge N dimension and C dimension to 1 dimension but still keep 4 dimensions to ensure the data can be used for training. Args: input_dict (Dict[str, paddle.Tensor]): input data dict. Returns: Dict[str, paddle.Tensor]: reshaped data dict. """out_dict={}forkeyininput_dict:input=input_dict[key]N,C,H,W=input.shapeout_dict[key]=paddle.reshape(input,[N*C,1,H,W])returnout_dictdefdereshape_input(input_dict:Dict[str,paddle.Tensor],C:int)->Dict[str,paddle.Tensor]:"""Dereshape input data for temporally Discriminator. Deeshape data from 1, N * C, H, W to N, C, W, H. Args: input_dict (Dict[str, paddle.Tensor]): input data dict. C (int): Channel of dereshape. Returns: Dict[str, paddle.Tensor]: dereshaped data dict. """forkeyininput_dict:input=input_dict[key]_,N,H,W=input.shapeifN<C:logger.warning(f"batch_size is smaller than {C}! Tempo needs at least {C} frames, input will be copied.")input_dict[key]=paddle.concat([input[:1]]*C,axis=1)else:N_new=int(N//C)input_dict[key]=paddle.reshape(input[:N_new*C],[-1,C,H,W])returninput_dict# predictdefsplit_data(data:np.ndarray,tile_ratio:int)->np.ndarray:"""Split a numpy image to tiles equally. Args: data (np.ndarray): The image to be Split. tile_ratio (int): How many tiles of one dim. Number of result tiles is tile_ratio * tile_ratio for a 2d image. Returns: np.ndarray: Tiles in [N,C,H,W] shape. """_,_,h,w=data.shapetile_h,tile_w=h//tile_ratio,w//tile_ratiotiles=[]foriinrange(tile_ratio):forjinrange(tile_ratio):tiles.append(data[:1,:,i*tile_h:i*tile_h+tile_h,j*tile_w:j*tile_w+tile_w,],)returnnp.concatenate(tiles,axis=0)defconcat_data(data:np.ndarray,tile_ratio:int)->np.ndarray:"""Concat numpy tiles to a image equally. Args: data (np.ndarray): The tiles to be upsplited. tile_ratio (int): How many tiles of one dim. Number of input tiles is tile_ratio * tile_ratio for 2d result. Returns: np.ndarray: Image in [H,W] shape. """_,_,tile_h,tile_w=data.shapeh,w=tile_h*tile_ratio,tile_w*tile_ratiodata_whole=np.ones([h,w],dtype=paddle.get_default_dtype())tile_idx=0foriinrange(tile_ratio):forjinrange(tile_ratio):data_whole[i*tile_h:i*tile_h+tile_h,j*tile_w:j*tile_w+tile_w,]=data[tile_idx][0]tile_idx+=1returndata_wholedefpredict_and_save_plot(output_dir:str,epoch_id:int,solver_gen:ppsci.solver.Solver,dataset_valid:np.ndarray,tile_ratio:int=1,):"""Predicting and plotting. Args: output_dir (str): Output dir path. epoch_id (int): Which epoch it is. solver_gen (ppsci.solver.Solver): Solver for predicting. dataset_valid (np.ndarray): Valid dataset. tile_ratio (int, optional): How many tiles of one dim. Defaults to 1. """dir_pred="predict/"os.makedirs(os.path.join(output_dir,dir_pred),exist_ok=True)start_idx=190density_low=dataset_valid["density_low"][start_idx:start_idx+3]density_high=dataset_valid["density_high"][start_idx:start_idx+3]# tiledensity_low=(split_data(density_low,tile_ratio)iftile_ratio!=1elsedensity_low)density_high=(split_data(density_high,tile_ratio)iftile_ratio!=1elsedensity_high)pred_dict=solver_gen.predict({"density_low":density_low,"density_high":density_high,},{"density_high":lambdaout:out["output_gen"]},batch_size=tile_ratio*tile_ratioiftile_ratio!=1else3,no_grad=False,)ifepoch_id==1:# plot interpolated input imageinput_img=np.expand_dims(dataset_valid["density_low"][start_idx],axis=0)input_img=paddle.to_tensor(input_img,dtype=paddle.get_default_dtype())input_img=F.interpolate(input_img,[input_img.shape[-2]*4,input_img.shape[-1]*4],mode="nearest",).numpy()Img.imsave(os.path.join(output_dir,dir_pred,"input.png"),np.squeeze(input_img),vmin=0.0,vmax=1.0,cmap="gray",)# plot target imageImg.imsave(os.path.join(output_dir,dir_pred,"target.png"),np.squeeze(dataset_valid["density_high"][start_idx]),vmin=0.0,vmax=1.0,cmap="gray",)# plot pred imagepred_img=(concat_data(pred_dict["density_high"].numpy(),tile_ratio)iftile_ratio!=1elsenp.squeeze(pred_dict["density_high"][0].numpy()))Img.imsave(os.path.join(output_dir,dir_pred,f"pred_epoch_{str(epoch_id)}.png"),pred_img,vmin=0.0,vmax=1.0,cmap="gray",)# evaluationdefevaluate_img(img_target:np.ndarray,img_pred:np.ndarray)->Tuple[float,float,float]:"""Evaluate two images. Args: img_target (np.ndarray): Target image. img_pred (np.ndarray): Image generated by prediction. Returns: Tuple[float, float, float]: MSE, PSNR, SSIM. """eval_mse=mean_squared_error(img_target,img_pred)eval_psnr=peak_signal_noise_ratio(img_target,img_pred)eval_ssim=structural_similarity(img_target,img_pred,data_range=1.0)returneval_mse,eval_psnr,eval_ssimdefget_image_array(img_path):returnnp.array(Image.open(img_path).convert("L"))classGenFuncs:"""All functions used for Generator, including functions of transform and loss. Args: weight_gen (List[float]): Weights of L1 loss. weight_gen_layer (List[float], optional): Weights of layers loss. Defaults to None. """def__init__(self,weight_gen:List[float],weight_gen_layer:List[float]=None)->None:self.weight_gen=weight_genself.weight_gen_layer=weight_gen_layerdeftransform_in(self,_in):ratio=2input_dict=reshape_input(_in)density_low=input_dict["density_low"]density_low_inp=interpolate(density_low,ratio,"nearest")return{"input_gen":density_low_inp}defloss_func_gen(self,output_dict:Dict,*args)->paddle.Tensor:"""Calculate loss of generator when use spatial discriminator. The loss consists of l1 loss, l2 loss and layer loss when use spatial discriminator. Notice that all item of loss is optional because weight of them might be 0. Args: output_dict (Dict): output dict of model. Returns: paddle.Tensor: Loss of generator. """# l1 lossloss_l1=F.l1_loss(output_dict["output_gen"],output_dict["density_high"],"mean")losses=loss_l1*self.weight_gen[0]# l2 lossloss_l2=F.mse_loss(output_dict["output_gen"],output_dict["density_high"],"mean")losses+=loss_l2*self.weight_gen[1]ifself.weight_gen_layerisnotNone:# disc(generator_out) lossout_disc_from_gen=output_dict["out_disc_from_gen"][-1]label_ones=paddle.ones_like(out_disc_from_gen)loss_gen=F.binary_cross_entropy_with_logits(out_disc_from_gen,label_ones,reduction="mean")losses+=loss_gen# layer losskey_list=list(output_dict.keys())# ["out0_layer0","out0_layer1","out0_layer2","out0_layer3","out_disc_from_target",# "out1_layer0","out1_layer1","out1_layer2","out1_layer3","out_disc_from_gen"]loss_layer=0foriinrange(1,len(self.weight_gen_layer)):# i = 0,1,2,3loss_layer+=(self.weight_gen_layer[i]*F.mse_loss(output_dict[key_list[i]],output_dict[key_list[5+i]],reduction="sum",)/2)losses+=loss_layer*self.weight_gen_layer[0]return{"output_gen":losses}defloss_func_gen_tempo(self,output_dict:Dict,*args)->paddle.Tensor:"""Calculate loss of generator when use temporal discriminator. The loss is cross entropy loss when use temporal discriminator. Args: output_dict (Dict): output dict of model. Returns: paddle.Tensor: Loss of generator. """out_disc_tempo_from_gen=output_dict["out_disc_tempo_from_gen"][-1]label_t_ones=paddle.ones_like(out_disc_tempo_from_gen)loss_gen_t=F.binary_cross_entropy_with_logits(out_disc_tempo_from_gen,label_t_ones,reduction="mean")losses=loss_gen_t*self.weight_gen[2]return{"out_disc_tempo_from_gen":losses}classDiscFuncs:"""All functions used for Discriminator and temporally Discriminator, including functions of transform and loss. Args: weight_disc (float): Weight of loss generated by the discriminator to judge the true target. """def__init__(self,weight_disc:float)->None:self.weight_disc=weight_discself.model_gen=Nonedeftransform_in(self,_in):ratio=2input_dict=reshape_input(_in)density_low=input_dict["density_low"]density_high_from_target=input_dict["density_high"]density_low_inp=interpolate(density_low,ratio,"nearest")density_high_from_gen=self.model_gen(input_dict)["output_gen"]density_high_from_gen.stop_gradient=Truedensity_input_from_target=paddle.concat([density_low_inp,density_high_from_target],axis=1)density_input_from_gen=paddle.concat([density_low_inp,density_high_from_gen],axis=1)return{"input_disc_from_target":density_input_from_target,"input_disc_from_gen":density_input_from_gen,}deftransform_in_tempo(self,_in):density_high_from_target=_in["density_high"]input_dict=reshape_input(_in)density_high_from_gen=self.model_gen(input_dict)["output_gen"]density_high_from_gen.stop_gradient=Trueinput_trans={"input_tempo_disc_from_target":density_high_from_target,"input_tempo_disc_from_gen":density_high_from_gen,}returndereshape_input(input_trans,3)defloss_func(self,output_dict,*args):out_disc_from_target=output_dict["out_disc_from_target"]out_disc_from_gen=output_dict["out_disc_from_gen"]label_ones=paddle.ones_like(out_disc_from_target)label_zeros=paddle.zeros_like(out_disc_from_gen)loss_disc_from_target=F.binary_cross_entropy_with_logits(out_disc_from_target,label_ones,reduction="mean")loss_disc_from_gen=F.binary_cross_entropy_with_logits(out_disc_from_gen,label_zeros,reduction="mean")losses=loss_disc_from_target*self.weight_disc+loss_disc_from_genreturn{"CE_loss":losses}defloss_func_tempo(self,output_dict,*args):out_disc_tempo_from_target=output_dict["out_disc_tempo_from_target"]out_disc_tempo_from_gen=output_dict["out_disc_tempo_from_gen"]label_ones=paddle.ones_like(out_disc_tempo_from_target)label_zeros=paddle.zeros_like(out_disc_tempo_from_gen)loss_disc_tempo_from_target=F.binary_cross_entropy_with_logits(out_disc_tempo_from_target,label_ones,reduction="mean")loss_disc_tempo_from_gen=F.binary_cross_entropy_with_logits(out_disc_tempo_from_gen,label_zeros,reduction="mean")losses=(loss_disc_tempo_from_target*self.weight_disc+loss_disc_tempo_from_gen)return{"CE_tempo_loss":losses}classDataFuncs:"""All functions used for data transform. Args: tile_ratio (int, optional): How many tiles of one dim. Defaults to 1. density_min (float, optional): Minimize density of one tile. Defaults to 0.02. max_turn (int, optional): Maximize turn of taking a tile from one image. Defaults to 20. """def__init__(self,tile_ratio:int=1,density_min:float=0.02,max_turn:int=20)->None:self.tile_ratio=tile_ratioself.density_min=density_minself.max_turn=max_turndeftransform(self,input_item:Dict[str,np.ndarray],label_item:Dict[str,np.ndarray],weight_item:Dict[str,np.ndarray],)->Tuple[Dict[str,np.ndarray],Dict[str,np.ndarray],Dict[str,np.ndarray]]:ifself.tile_ratio==1:returninput_item,label_item,weight_itemfor_inrange(self.max_turn):rand_ratio=np.random.rand()density_low=self.cut_data(input_item["density_low"],rand_ratio)density_high=self.cut_data(input_item["density_high"],rand_ratio)ifself.is_valid_tile(density_low):breakinput_item["density_low"]=density_lowinput_item["density_high"]=density_highreturninput_item,label_item,weight_itemdefcut_data(self,data:np.ndarray,rand_ratio:float)->paddle.Tensor:# data: C,H,W_,H,W=data.shapeifH%self.tile_ratio!=0orW%self.tile_ratio!=0:exit(f"ERROR: input images cannot be divided into {self.tile_ratio} parts evenly!")tile_shape=[H//self.tile_ratio,W//self.tile_ratio]rand_shape=np.floor(rand_ratio*(np.array([H,W])-np.array(tile_shape)))start=[int(rand_shape[0]),int(rand_shape[1])]end=[int(rand_shape[0]+tile_shape[0]),int(rand_shape[1]+tile_shape[1])]data=paddle.slice(paddle.to_tensor(data),axes=[-2,-1],starts=start,ends=end)returndatadefis_valid_tile(self,tile:paddle.Tensor):img_density=tile[0].sum()returnimg_density>=(self.density_min*tile.shape[0]*tile.shape[1]*tile.shape[2])
# Copyright (c) 2023 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.from__future__importannotationsfromtypingimportDictfromtypingimportListfromtypingimportTupleimportpaddleimportpaddle.nnasnnfromppsci.archimportactivationasact_modfromppsci.archimportbaseclassConv2DBlock(nn.Layer):def__init__(self,in_channel,out_channel,kernel_size,stride,use_bn,act,mean,std,value,):super().__init__()weight_attr=paddle.ParamAttr(initializer=nn.initializer.Normal(mean=mean,std=std))bias_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=value))self.conv_2d=nn.Conv2D(in_channel,out_channel,kernel_size,stride,padding="SAME",weight_attr=weight_attr,bias_attr=bias_attr,)self.bn=nn.BatchNorm2D(out_channel)ifuse_bnelseNoneself.act=act_mod.get_activation(act)ifactelseNonedefforward(self,x):y=xy=self.conv_2d(y)ifself.bn:y=self.bn(y)ifself.act:y=self.act(y)returnyclassVariantResBlock(nn.Layer):def__init__(self,in_channel,out_channels,kernel_sizes,strides,use_bns,acts,mean,std,value,):super().__init__()self.conv_2d_0=Conv2DBlock(in_channel=in_channel,out_channel=out_channels[0],kernel_size=kernel_sizes[0],stride=strides[0],use_bn=use_bns[0],act=acts[0],mean=mean,std=std,value=value,)self.conv_2d_1=Conv2DBlock(in_channel=out_channels[0],out_channel=out_channels[1],kernel_size=kernel_sizes[1],stride=strides[1],use_bn=use_bns[1],act=acts[1],mean=mean,std=std,value=value,)self.conv_2d_2=Conv2DBlock(in_channel=in_channel,out_channel=out_channels[2],kernel_size=kernel_sizes[2],stride=strides[2],use_bn=use_bns[2],act=acts[2],mean=mean,std=std,value=value,)self.act=act_mod.get_activation("relu")defforward(self,x):y=xy=self.conv_2d_0(y)y=self.conv_2d_1(y)short=self.conv_2d_2(x)y=paddle.add(y,short)y=self.act(y)returnyclassFCBlock(nn.Layer):def__init__(self,in_channel,act,mean,std,value):super().__init__()self.flatten=nn.Flatten()weight_attr=paddle.ParamAttr(initializer=nn.initializer.Normal(mean=mean,std=std))bias_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=value))self.linear=nn.Linear(in_channel,1,weight_attr=weight_attr,bias_attr=bias_attr,)self.act=act_mod.get_activation(act)ifactelseNonedefforward(self,x):y=xy=self.flatten(y)y=self.linear(y)ifself.act:y=self.act(y)returnyclassGenerator(base.Arch):"""Generator Net of GAN. Attention, the net using a kind of variant of ResBlock which is unique to "tempoGAN" example but not an open source network. Args: input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). in_channel (int): Number of input channels of the first conv layer. out_channels_tuple (Tuple[Tuple[int, ...], ...]): Number of output channels of all conv layers, such as [[out_res0_conv0, out_res0_conv1], [out_res1_conv0, out_res1_conv1]] kernel_sizes_tuple (Tuple[Tuple[int, ...], ...]): Number of kernel_size of all conv layers, such as [[kernel_size_res0_conv0, kernel_size_res0_conv1], [kernel_size_res1_conv0, kernel_size_res1_conv1]] strides_tuple (Tuple[Tuple[int, ...], ...]): Number of stride of all conv layers, such as [[stride_res0_conv0, stride_res0_conv1], [stride_res1_conv0, stride_res1_conv1]] use_bns_tuple (Tuple[Tuple[bool, ...], ...]): Whether to use the batch_norm layer after each conv layer. acts_tuple (Tuple[Tuple[str, ...], ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, such as [[act_res0_conv0, act_res0_conv1], [act_res1_conv0, act_res1_conv1]] Examples: >>> import ppsci >>> in_channel = 1 >>> rb_channel0 = (2, 8, 8) >>> rb_channel1 = (128, 128, 128) >>> rb_channel2 = (32, 8, 8) >>> rb_channel3 = (2, 1, 1) >>> out_channels_tuple = (rb_channel0, rb_channel1, rb_channel2, rb_channel3) >>> kernel_sizes_tuple = (((5, 5), ) * 2 + ((1, 1), ), ) * 4 >>> strides_tuple = ((1, 1, 1), ) * 4 >>> use_bns_tuple = ((True, True, True), ) * 3 + ((False, False, False), ) >>> acts_tuple = (("relu", None, None), ) * 4 >>> model = ppsci.arch.Generator(("in",), ("out",), in_channel, out_channels_tuple, kernel_sizes_tuple, strides_tuple, use_bns_tuple, acts_tuple) >>> batch_size = 4 >>> height = 64 >>> width = 64 >>> input_data = paddle.randn([batch_size, in_channel, height, width]) >>> input_dict = {'in': input_data} >>> output_data = model(input_dict) >>> print(output_data['out'].shape) [4, 1, 64, 64] """def__init__(self,input_keys:Tuple[str,...],output_keys:Tuple[str,...],in_channel:int,out_channels_tuple:Tuple[Tuple[int,...],...],kernel_sizes_tuple:Tuple[Tuple[int,...],...],strides_tuple:Tuple[Tuple[int,...],...],use_bns_tuple:Tuple[Tuple[bool,...],...],acts_tuple:Tuple[Tuple[str,...],...],):super().__init__()self.input_keys=input_keysself.output_keys=output_keysself.in_channel=in_channelself.out_channels_tuple=out_channels_tupleself.kernel_sizes_tuple=kernel_sizes_tupleself.strides_tuple=strides_tupleself.use_bns_tuple=use_bns_tupleself.acts_tuple=acts_tupleself.init_blocks()definit_blocks(self):blocks_list=[]foriinrange(len(self.out_channels_tuple)):in_channel=(self.in_channelifi==0elseself.out_channels_tuple[i-1][-1])blocks_list.append(VariantResBlock(in_channel=in_channel,out_channels=self.out_channels_tuple[i],kernel_sizes=self.kernel_sizes_tuple[i],strides=self.strides_tuple[i],use_bns=self.use_bns_tuple[i],acts=self.acts_tuple[i],mean=0.0,std=0.04,value=0.1,))self.blocks=nn.LayerList(blocks_list)defforward_tensor(self,x):y=xforblockinself.blocks:y=block(y)returnydefforward(self,x):ifself._input_transformisnotNone:x=self._input_transform(x)y=self.concat_to_tensor(x,self.input_keys,axis=-1)y=self.forward_tensor(y)y=self.split_to_dict(y,self.output_keys,axis=-1)ifself._output_transformisnotNone:y=self._output_transform(x,y)returnyclassDiscriminator(base.Arch):"""Discriminator Net of GAN. Args: input_keys (Tuple[str, ...]): Name of input keys, such as ("input1", "input2"). output_keys (Tuple[str, ...]): Name of output keys, such as ("output1", "output2"). in_channel (int): Number of input channels of the first conv layer. out_channels (Tuple[int, ...]): Number of output channels of all conv layers, such as (out_conv0, out_conv1, out_conv2). fc_channel (int): Number of input features of linear layer. Number of output features of the layer is set to 1 in this Net to construct a fully_connected layer. kernel_sizes (Tuple[int, ...]): Number of kernel_size of all conv layers, such as (kernel_size_conv0, kernel_size_conv1, kernel_size_conv2). strides (Tuple[int, ...]): Number of stride of all conv layers, such as (stride_conv0, stride_conv1, stride_conv2). use_bns (Tuple[bool, ...]): Whether to use the batch_norm layer after each conv layer. acts (Tuple[str, ...]): Whether to use the activation layer after each conv layer. If so, witch activation to use, such as (act_conv0, act_conv1, act_conv2). Examples: >>> import ppsci >>> in_channel = 2 >>> in_channel_tempo = 3 >>> out_channels = (32, 64, 128, 256) >>> fc_channel = 65536 >>> kernel_sizes = ((4, 4), (4, 4), (4, 4), (4, 4)) >>> strides = (2, 2, 2, 1) >>> use_bns = (False, True, True, True) >>> acts = ("leaky_relu", "leaky_relu", "leaky_relu", "leaky_relu", None) >>> output_keys_disc = ("out_1", "out_2", "out_3", "out_4", "out_5", "out_6", "out_7", "out_8", "out_9", "out_10") >>> model = ppsci.arch.Discriminator(("in_1","in_2"), output_keys_disc, in_channel, out_channels, fc_channel, kernel_sizes, strides, use_bns, acts) >>> input_data = [paddle.to_tensor(paddle.randn([1, in_channel, 128, 128])),paddle.to_tensor(paddle.randn([1, in_channel, 128, 128]))] >>> input_dict = {"in_1": input_data[0],"in_2": input_data[1]} >>> out_dict = model(input_dict) >>> for k, v in out_dict.items(): ... print(k, v.shape) out_1 [1, 32, 64, 64] out_2 [1, 64, 32, 32] out_3 [1, 128, 16, 16] out_4 [1, 256, 16, 16] out_5 [1, 1] out_6 [1, 32, 64, 64] out_7 [1, 64, 32, 32] out_8 [1, 128, 16, 16] out_9 [1, 256, 16, 16] out_10 [1, 1] """def__init__(self,input_keys:Tuple[str,...],output_keys:Tuple[str,...],in_channel:int,out_channels:Tuple[int,...],fc_channel:int,kernel_sizes:Tuple[int,...],strides:Tuple[int,...],use_bns:Tuple[bool,...],acts:Tuple[str,...],):super().__init__()self.input_keys=input_keysself.output_keys=output_keysself.in_channel=in_channelself.out_channels=out_channelsself.fc_channel=fc_channelself.kernel_sizes=kernel_sizesself.strides=stridesself.use_bns=use_bnsself.acts=actsself.init_layers()definit_layers(self):layers_list=[]foriinrange(len(self.out_channels)):in_channel=self.in_channelifi==0elseself.out_channels[i-1]layers_list.append(Conv2DBlock(in_channel=in_channel,out_channel=self.out_channels[i],kernel_size=self.kernel_sizes[i],stride=self.strides[i],use_bn=self.use_bns[i],act=self.acts[i],mean=0.0,std=0.04,value=0.1,))layers_list.append(FCBlock(self.fc_channel,self.acts[4],mean=0.0,std=0.04,value=0.1))self.layers=nn.LayerList(layers_list)defforward_tensor(self,x):y=xy_list=[]forlayerinself.layers:y=layer(y)y_list.append(y)returny_list# y_conv1, y_conv2, y_conv3, y_conv4, y_fc(y_out)defforward(self,x):ifself._input_transformisnotNone:x=self._input_transform(x)y_list=[]# y1_conv1, y1_conv2, y1_conv3, y1_conv4, y1_fc, y2_conv1, y2_conv2, y2_conv3, y2_conv4, y2_fcforkinx:y_list.extend(self.forward_tensor(x[k]))y=self.split_to_dict(y_list,self.output_keys)ifself._output_transformisnotNone:y=self._output_transform(x,y)returny@staticmethoddefsplit_to_dict(data_list:List[paddle.Tensor],keys:Tuple[str,...])->Dict[str,paddle.Tensor]:"""Overwrite of split_to_dict() method belongs to Class base.Arch. Reason for overwriting is there is no concat_to_tensor() method called in "tempoGAN" example. That is because input in "tempoGAN" example is not in a regular format, but a format like: { "input1": paddle.concat([in1, in2], axis=1), "input2": paddle.concat([in1, in3], axis=1), } Args: data_list (List[paddle.Tensor]): The data to be split. It should be a list of tensor(s), but not a paddle.Tensor. keys (Tuple[str, ...]): Keys of outputs. Returns: Dict[str, paddle.Tensor]: Dict with split data. """iflen(keys)==1:return{keys[0]:data_list[0]}return{key:data_list[i]fori,keyinenumerate(keys)}
After using mixed precision training, evaluate MSE, PSNR, SSIM between target on test set. The values of evaluation indicators are:
MSE
PSNR
SSIM
4.21e-5
47.19
0.9974
The input of a fluid super-resolution example, model prediction result, and result directly generated by open source code package mantaflow in Dataset Introduction are as follows. The model prediction result is basically consistent with the generated target result.
Input low-density fluid
High-density fluid obtained by inference after mixed precision training
