├── .gitignore ├── LICENSE ├── README.md ├── ReadmeFiles ├── DataStruct.png ├── arch.png ├── circle1.png ├── circle2.png ├── rhombus1.png ├── rhombus2.png ├── square1.png ├── square2.png ├── triangleB1.png ├── triangleB2.png ├── triangleF1.png └── triangleF2.png ├── scripts ├── visualize_data.py └── visualize_training_logs.py ├── setup.py └── src └── deepcfd ├── __init__.py ├── __main__.py ├── functions.py ├── models ├── AutoEncoder.py ├── AutoEncoderEx.py ├── UNet.py ├── UNetEx.py ├── UNetExAvg.py ├── UNetExMod.py └── __init__.py ├── pytorchtools.py └── train_functions.py /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | .idea/ 3 | 4 | *.pkl 5 | 6 | .ipynb_checkpoints/ 7 | 8 | __pycache__/ 9 | 10 | *.txt 11 | 12 | Results/ 13 | 14 | BestRun/ 15 | 16 | _Results/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Mateus Dias Ribeiro 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepCFD 2 | 3 | Computational Fluid Dynamics (CFD) simulation by the numerical solution of the Navier-Stokes equations is an essential tool in a wide range of applications from engineering design to climate modeling. However, the computational cost and memory demand required by CFD codes may become very high for flows of practical interest, such as in aerodynamic shape optimization. This expense is associated with the complexity of the fluid flow governing equations, which include non-linear partial derivative terms that are of difficult solution, leading to long computational times and limiting the number of hypotheses that can be tested during the process of iterative design. Therefore, we propose DeepCFD: a convolutional neural network (CNN) based model that efficiently approximates solutions for the problem of non-uniform steady laminar flows. The proposed model is able to learn complete solutions of the Navier-Stokes equations, for both velocity and pressure fields, directly from ground-truth data generated using a state-of-the-art CFD code. Using DeepCFD, we found a speedup of up to 3 orders of magnitude compared to the standard CFD approach at a cost of low error rates. 4 | 5 | Paper: https://arxiv.org/abs/2004.08826 6 | 7 | --- 8 | 9 | ## Installation 10 | 11 | The module can be installed with: 12 | 13 | ``` 14 | pip3 install git+https://github.com/mdribeiro/DeepCFD.git@master 15 | 16 | ``` 17 | 18 | ## Usage 19 | 20 | ``` 21 | Usage: python3 -m deepcfd [OPTIONS] 22 | 23 | Options: 24 | -d, --device TEXT device: 'cpu', 'cuda', 'cuda:0', 'cuda:0,cuda:1' (default: cuda if available) 25 | -n, --net TEXT network architecture: UNetEx or AutoEncoder (default: UNetEx) 26 | -mi, --mmodel-input PATH input dataset with sdf1,flow-region and sdf2 fields (default: dataX.pkl) 27 | -mo, --model-output PATH output dataset with Ux,Uy and p (default: dataY.pkl) 28 | -o, --output PATH model output (default: mymodel.pt) 29 | -k, --kernel-size INT kernel size (default: 5) 30 | -f, --filters TEXT filter size (default: 8,16,32,32) 31 | -l, --learning-rate FLOAT learning rage (default: 0.001) 32 | -e, --epochs INT number of epochs (default: 1000) 33 | -b, --batch-size INT training batch size (default: 32) 34 | -p, --patience INT number of epochs for early stopping (default: 300) 35 | -v, --visualize flag for visualizing ground-truth vs prediction plots (default: False) 36 | 37 | 38 | Example: 39 | 40 | python3 -m deepcfd \ 41 | --net UNetEx \ 42 | --model-input DeepCFD/$dataX.pkl \ 43 | --model-output DeepCFD/dataY.pkl \ 44 | --output DeepCFD/${name}.pt \ 45 | --kernel-size 5 \ 46 | --filters 8,16,32,32 \ 47 | --epochs 2000 \ 48 | --batch-size 32 > log.deepcfd 49 | ``` 50 | 51 | ## Dataset and Code 52 | 53 | A toy dataset and the code for this project can be downloaded using the following https://zenodo.org/record/3666056/files/DeepCFD.zip?download=1 54 | 55 | The folder includes the files dataX and dataY, in which the first file provides the input information on the geometry of 981 channel flow samples, whereas the dataY file provides their ground-truth CFD solution for the velocity (Ux and Uy) and thepressure (p) fields using the simpleFOAM solver. Figure 1 describes the structure of each of these files in detail: 56 | 57 | ![CFDAI](./ReadmeFiles/DataStruct.png) 58 | > Figure 1. DeepCFD toy dataset structure. 59 | 60 | Both dataX/dataY have the same dimensions (Ns, Nc, Nx, Ny), in which the first axis is the number of samples (Ns), the second axis is the number of channels (Nc), and third and fourth axes are the number of elements in x and y (Nx and Ny). Regarding the input dataX, the first channel is the SDF calculated from the obstacle's surface, the second channel is the multi-label flow region channel, and the third channel is the SDF from the top/bottom surfaces. For the output dataY file, the first channel is the Ux horizontal velocity component, the second channel is the Uy vertical velocity component, and the third channel is the pressure field. 61 | 62 | An example of how to train the DeepCFD model using the settings described in the paper is provided in the "DeepCFD.py" script. A few useful functions are provided in the "functions.py" file, such as a plotting function to visualize the outcome of the model. Moreover, templates with all networks investigated in this study can be found in the folder "Models", including both "AutoEncoder" and "UNet" architecture types with one or multiple decoders. 63 | 64 | ## Flow Visualization Plots - OpenFOAM vs DeepCFD 65 | 66 | ![Circle1](./ReadmeFiles/circle1.png) 67 | > Figure 2. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around circle based shape 1. 68 | 69 | ![Circle2](./ReadmeFiles/circle2.png) 70 | > Figure 3. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around circle based shape 2. 71 | 72 | ![Square1](./ReadmeFiles/square1.png) 73 | > Figure 4. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around square based shape 1. 74 | 75 | ![Square2](./ReadmeFiles/square2.png) 76 | > Figure 5. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around square based shape 2. 77 | 78 | 79 | ![Rhombus1](./ReadmeFiles/rhombus1.png) 80 | > Figure 6. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around rhombus based shape 1. 81 | 82 | ![Rhombus2](./ReadmeFiles/rhombus2.png) 83 | > Figure 7. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around rhombus based shape 2. 84 | 85 | 86 | ![TriangleF1](./ReadmeFiles/triangleF1.png) 87 | > Figure 8. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around forward-facing triangle based shape 1. 88 | 89 | ![TriangleF2](./ReadmeFiles/triangleF2.png) 90 | > Figure 9. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around forward-facing triangle based shape 2. 91 | 92 | ![TriangleB1](./ReadmeFiles/triangleB1.png) 93 | > Figure 10. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around backward-facing triangle based shape 1. 94 | 95 | ![TriangleB2](./ReadmeFiles/triangleB2.png) 96 | > Figure 11. Comparison between ground-truth CFD (simpleFOAM) and DeepCFD prediction, showing both velocity components, and pressure fields, as well as absolute error in flow around backward-facing triangle based shape 2. 97 | 98 | ## DeepCFD architecture 99 | 100 | Figure 12 shows a schematic representation of the DeepCFD architecture: 101 | 102 | ![arch](./ReadmeFiles/arch.png) 103 | > Figure 12. DeepCFD U-Net Architecture. 104 | 105 | Further hyper-parameters used are provided in Table 1: 106 | 107 | > Table 1. DeepCFD hyper-parameters. 108 | 109 | | Parameter | Value | 110 | | ------------- |--------------:| 111 | | Learning rate | 1e-3 | 112 | | Kernel size | 5 | 113 | | Filters | [8,16,32,32] | 114 | | Batch size | 64 | 115 | | Weight decay | 0.005 | 116 | | Batch norm | off | 117 | | Weight norm | off | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | -------------------------------------------------------------------------------- /ReadmeFiles/DataStruct.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/DataStruct.png -------------------------------------------------------------------------------- /ReadmeFiles/arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/arch.png -------------------------------------------------------------------------------- /ReadmeFiles/circle1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/circle1.png -------------------------------------------------------------------------------- /ReadmeFiles/circle2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/circle2.png -------------------------------------------------------------------------------- /ReadmeFiles/rhombus1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/rhombus1.png -------------------------------------------------------------------------------- /ReadmeFiles/rhombus2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/rhombus2.png -------------------------------------------------------------------------------- /ReadmeFiles/square1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/square1.png -------------------------------------------------------------------------------- /ReadmeFiles/square2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/square2.png -------------------------------------------------------------------------------- /ReadmeFiles/triangleB1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/triangleB1.png -------------------------------------------------------------------------------- /ReadmeFiles/triangleB2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/triangleB2.png -------------------------------------------------------------------------------- /ReadmeFiles/triangleF1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/triangleF1.png -------------------------------------------------------------------------------- /ReadmeFiles/triangleF2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/ReadmeFiles/triangleF2.png -------------------------------------------------------------------------------- /scripts/visualize_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import time 4 | import deepcfd 5 | import torch 6 | import pickle 7 | import matplotlib 8 | import numpy as np 9 | from deepcfd.models.UNetEx import UNetEx 10 | from matplotlib import pyplot as plt 11 | 12 | matplotlib.use('tkagg') 13 | 14 | model_filename="checkpoint.pt" 15 | data_x_filename="dataX.pkl" 16 | data_y_filename="dataY.pkl" 17 | 18 | index = 200 19 | 20 | kernel_size = 5 21 | filters = [8, 16, 32, 32] 22 | bn = False 23 | wn = False 24 | 25 | last_mtime = 0 26 | last_saved = 1 27 | current_epoch = 1 28 | extent=[0,256,0,128] 29 | 30 | min_u_x = 0 31 | max_u_x = 0.15 32 | min_u_y = -0.045 33 | max_u_y = 0.045 34 | min_u_x_error = 0 35 | max_u_x_error = 0.0185 36 | min_u_y_error = 0 37 | max_u_y_error = 0.0085 38 | min_p = 0 39 | max_p = 0.015 40 | min_p_error = 0 41 | max_p_error = 0.0075 42 | 43 | model = UNetEx( 44 | 3, 45 | 3, 46 | filters=filters, 47 | kernel_size=kernel_size, 48 | batch_norm=bn, 49 | weight_norm=wn 50 | ) 51 | 52 | x = pickle.load(open(data_x_filename, "rb")) 53 | y = pickle.load(open(data_y_filename, "rb")) 54 | 55 | x = torch.FloatTensor(x) 56 | y = torch.FloatTensor(y) 57 | 58 | truth = y[index:(index+1)].cpu().detach().numpy() 59 | inputs = x[index:(index+1)].cpu().detach().numpy() 60 | 61 | plt.figure() 62 | fig = plt.gcf() 63 | fig.set_size_inches(15, 10) 64 | fig.canvas.manager.window.wm_attributes('-topmost', 0) 65 | 66 | plt.ion() 67 | plt.show() 68 | 69 | def update_plot(): 70 | 71 | # Wait until the file is completely written 72 | while True: 73 | try: 74 | with open(model_filename, "r") as f: 75 | break 76 | except IOError: 77 | time.sleep(1) 78 | 79 | state_dict = torch.load(model_filename) 80 | model.load_state_dict(state_dict) 81 | out = model(x[index:(index+1)]).detach().numpy() 82 | error = abs(out - truth) 83 | 84 | fig.suptitle('Best at epoch ' + str(last_saved) + ' (current: ' + str(current_epoch) + ')') 85 | 86 | plt.subplot(3, 3, 1) 87 | plt.ylabel('Ux [m/s]', fontsize=18) 88 | plt.title('CFD', fontsize=18) 89 | plt.imshow(truth[0, 0, :, :], cmap='jet', vmin = min_u_x, vmax = max_u_x, origin='lower', extent=extent) 90 | plt.colorbar(orientation='horizontal') 91 | plt.subplot(3, 3, 2) 92 | plt.title('CNN', fontsize=18) 93 | plt.imshow(out[0, 0, :, :], cmap='jet', vmin = min_u_x, vmax = max_u_x, origin='lower', extent=extent) 94 | plt.colorbar(orientation='horizontal') 95 | plt.subplot(3, 3, 3) 96 | plt.title('error', fontsize=18) 97 | plt.imshow(error[0, 0, :, :], cmap='jet', vmin = min_u_x_error, vmax = max_u_x_error, origin='lower', extent=extent) 98 | plt.colorbar(orientation='horizontal') 99 | 100 | plt.subplot(3, 3, 4) 101 | plt.ylabel('Uy [m/s]', fontsize=18) 102 | plt.imshow(truth[0, 1, :, :], cmap='jet', vmin = min_u_y, vmax = max_u_y, origin='lower', extent=extent) 103 | plt.colorbar(orientation='horizontal') 104 | plt.subplot(3, 3, 5) 105 | plt.imshow(out[0, 1, :, :], cmap='jet', vmin = min_u_y, vmax = max_u_y, origin='lower', extent=extent) 106 | plt.colorbar(orientation='horizontal') 107 | plt.subplot(3, 3, 6) 108 | plt.imshow(error[0, 1, :, :], cmap='jet', vmin = min_u_y_error, vmax = max_u_y_error, origin='lower', extent=extent) 109 | plt.colorbar(orientation='horizontal') 110 | 111 | plt.subplot(3, 3, 7) 112 | plt.ylabel('p [m2/s2]', fontsize=18) 113 | plt.imshow(truth[0, 2, :, :], cmap='jet', vmin = min_p, vmax = max_p, origin='lower', extent=extent) 114 | plt.colorbar(orientation='horizontal') 115 | plt.subplot(3, 3, 8) 116 | plt.imshow(out[0, 2, :, :], cmap='jet', vmin = min_p, vmax = max_p, origin='lower', extent=extent) 117 | plt.colorbar(orientation='horizontal') 118 | plt.subplot(3, 3, 9) 119 | plt.imshow(error[0, 2, :, :], cmap='jet', vmin = min_p_error, vmax = max_p_error, origin='lower', extent=extent) 120 | plt.colorbar(orientation='horizontal') 121 | 122 | plt.draw() 123 | plt.pause(5) 124 | #plt.show() 125 | 126 | last_saved = current_epoch 127 | 128 | while True: 129 | saved_epochs = [] 130 | with open('log.deepcfd2', 'r') as f: 131 | lines = f.readlines() 132 | for line in reversed(lines): 133 | if line.startswith('Epoch #'): 134 | current_epoch = int(line.split('Epoch #')[1]) 135 | break 136 | for line in reversed(lines): 137 | if line.startswith('Epoch #'): 138 | last_saved = int(line.split('Epoch #')[1]) 139 | elif line.startswith('Model saved!'): 140 | last_saved = last_saved - 1 141 | break 142 | 143 | update_plot() 144 | 145 | time.sleep(5) 146 | -------------------------------------------------------------------------------- /scripts/visualize_training_logs.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | 3 | log_file = "log.deepcfd" 4 | initial_epoch = 20 5 | 6 | train_loss = [] 7 | train_total_mse = [] 8 | train_ux_mse = [] 9 | train_uy_mse = [] 10 | train_p_mse = [] 11 | val_loss = [] 12 | val_total_mse = [] 13 | val_ux_mse = [] 14 | val_uy_mse = [] 15 | val_p_mse = [] 16 | 17 | with open(log_file, 'r') as f: 18 | for line in f: 19 | if line.strip().startswith('Train Loss'): 20 | train_loss.append(float(line.split()[-1])) 21 | elif line.strip().startswith('Train Total MSE'): 22 | train_total_mse.append(float(line.split()[-1])) 23 | elif line.strip().startswith('Train Ux MSE'): 24 | train_ux_mse.append(float(line.split()[-1])) 25 | elif line.strip().startswith('Train Uy MSE'): 26 | train_uy_mse.append(float(line.split()[-1])) 27 | elif line.strip().startswith('Train p MSE'): 28 | train_p_mse.append(float(line.split()[-1])) 29 | elif line.strip().startswith('Validation Loss'): 30 | val_loss.append(float(line.split()[-1])) 31 | elif line.strip().startswith('Validation Total MSE'): 32 | val_total_mse.append(float(line.split()[-1])) 33 | elif line.strip().startswith('Validation Ux MSE'): 34 | val_ux_mse.append(float(line.split()[-1])) 35 | elif line.strip().startswith('Validation Uy MSE'): 36 | val_uy_mse.append(float(line.split()[-1])) 37 | elif line.strip().startswith('Validation p MSE'): 38 | val_p_mse.append(float(line.split()[-1])) 39 | 40 | epochs = list(range(1, len(train_loss) + 1)) 41 | 42 | fig, axs = plt.subplots(ncols=2, sharex=True, figsize=(14, 8)) 43 | 44 | axs[0].plot(epochs[initial_epoch:], train_loss[initial_epoch:], label='Train Loss') 45 | axs[0].plot(epochs[initial_epoch:], val_loss[initial_epoch:], label='Validation Loss') 46 | axs[0].set_xlabel('Epoch') 47 | axs[0].set_ylabel('Loss') 48 | axs[0].legend() 49 | 50 | axs[1].plot(epochs[initial_epoch:], train_total_mse[initial_epoch:], label='Train Total MSE') 51 | axs[1].plot(epochs[initial_epoch:], train_ux_mse[initial_epoch:], label='Train Ux MSE') 52 | axs[1].plot(epochs[initial_epoch:], train_uy_mse[initial_epoch:], label='Train Uy MSE') 53 | axs[1].plot(epochs[initial_epoch:], train_p_mse[initial_epoch:], label='Train p MSE') 54 | axs[1].plot(epochs[initial_epoch:], val_total_mse[initial_epoch:], label='Validation Total MSE') 55 | axs[1].plot(epochs[initial_epoch:], val_ux_mse[initial_epoch:], label='Validation Ux MSE') 56 | axs[1].plot(epochs[initial_epoch:], val_uy_mse[initial_epoch:], label='Validation Uy MSE') 57 | axs[1].plot(epochs[initial_epoch:], val_p_mse[initial_epoch:], label='Validation p MSE') 58 | axs[1].set_xlabel('Epoch') 59 | axs[1].set_ylabel('MSE') 60 | axs[1].legend() 61 | 62 | plt.show() 63 | 64 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | with open("README.md", "r") as fh: 4 | long_description = fh.read() 5 | 6 | setup( 7 | name='deepcfd', 8 | version='2.0.0', 9 | description='DeepCFD: Efficient Steady-State Laminar Flow' 10 | 'Approximation with Deep Convolutional Neural Networks', 11 | author='Mateus Dias Ribeiro, Carlos Peña Monferrer', 12 | author_email='mateusdiasbr@gmail.com', 13 | long_description=long_description, 14 | long_description_content_type='text/markdown', 15 | url='https://github.com/mdribeiro/DeepCFD', 16 | packages=['deepcfd', 'deepcfd.models'], 17 | package_dir={'': 'src'}, 18 | classifiers=[ 19 | "Programming Language :: Python :: 3", 20 | "License :: OSI Approved :: MIT License", 21 | "Operating System :: OS Independent", 22 | ], 23 | install_requires=[ 24 | "torch==2.0.0", 25 | "torchvision==0.15.1", 26 | "matplotlib>=3.0.0,<=3.7.1" 27 | ], 28 | ) 29 | -------------------------------------------------------------------------------- /src/deepcfd/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/src/deepcfd/__init__.py -------------------------------------------------------------------------------- /src/deepcfd/__main__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import torch 4 | import pickle 5 | import random 6 | import getopt 7 | import sys 8 | from .train_functions import * 9 | from .functions import * 10 | import torch.optim as optim 11 | from torch.utils.data import TensorDataset 12 | from torch.autograd import Variable 13 | 14 | 15 | def parseOpts(argv): 16 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 17 | net= "UNetEx" 18 | kernel_size = 5 19 | filters = [8, 16, 32, 32] 20 | model_input = "dataX.pkl" 21 | model_output = "dataY.pkl" 22 | output = "mymodel.pt" 23 | learning_rate = 0.001 24 | epochs = 2000 25 | batch_size = 32 26 | patience = 500 27 | visualize = False 28 | 29 | try: 30 | opts, args = getopt.getopt( 31 | argv,"hd:n:mi:mo:o:k:f:l:e:b:p:v", 32 | [ 33 | "device=", 34 | "net=", 35 | "model-input=", 36 | "model-output=", 37 | "output=", 38 | "kernel-size=", 39 | "filters=", 40 | "learning-rate=", 41 | "epochs=", 42 | "batch-size=", 43 | "patience=", 44 | "visualize" 45 | ] 46 | ) 47 | except getopt.GetoptError as e: 48 | print(e) 49 | print("python -m deepcfd --help") 50 | sys.exit(2) 51 | 52 | for opt, arg in opts: 53 | if opt == '-h' or opt == '--help': 54 | print("deepcfd " 55 | "\n -d device: 'cpu', 'cuda', 'cuda:0', 'cuda:0,cuda:n', (default: cuda if available)" 56 | "\n -n network architecture: UNet, UNetEx or " 57 | "AutoEncoder (default: UNetEx)" 58 | "\n -mi input dataset with sdf1," 59 | "flow-region and sdf2 fields (default: dataX.pkl)" 60 | "\n -mo output dataset with Ux," 61 | "Uy and p (default: dataY.pkl)" 62 | "\n -o model output (default: mymodel.pt)" 63 | "\n -k kernel size (default: 5)" 64 | "\n -f filter sizes (default: 8,16,32,32)" 65 | "\n -l learning rate (default: 0.001)" 66 | "\n -e number of epochs (default: 1000)" 67 | "\n -b training batch size (default: 32)" 68 | "\n -p number of epochs for early stopping (default: 300)" 69 | "\n -v flag for visualizing ground-truth vs prediction plots (default: False)\n" 70 | ) 71 | sys.exit() 72 | elif opt in ("-d", "--device"): 73 | if (arg == "cpu" or arg.startswith("cuda")): 74 | device = arg; 75 | else: 76 | print("Unkown device " + str(arg) + ", only 'cpu', 'cuda'" 77 | "'cuda:index', or comma-separated list of 'cuda:index'" 78 | "are supported") 79 | exit(0) 80 | elif opt in ("-n", "--net"): 81 | if (arg == "UNet"): 82 | from .models.UNet import UNet 83 | net = UNet 84 | elif (arg == "UNetEx"): 85 | from .models.UNetEx import UNetEx 86 | net = UNetEx 87 | elif (arg == "AutoEncoder"): 88 | from .models.AutoEncoder import AutoEncoder 89 | net = AutoEncoder 90 | else: 91 | print("Unkown network " + str(arg) + ", only UNet, UNetEx" 92 | "and AutoEncoder are supported") 93 | exit(0) 94 | elif opt in ("-mi", "--model-input"): 95 | model_input = arg 96 | elif opt in ("-mo", "--model-output"): 97 | model_output = arg 98 | elif opt in ("-k", "--kernel-size"): 99 | kernel_size = int(arg) 100 | elif opt in ("-f", "--filters"): 101 | filters = [int(x) for x in arg.split(',')] 102 | elif opt in ("-l", "--learning-rate"): 103 | learning_rate = float(arg) 104 | elif opt in ("-e", "--epochs"): 105 | epochs = int(arg) 106 | elif opt in ("-b", "--batch-size"): 107 | batch_size = int(arg) 108 | elif opt in ("-o", "--output"): 109 | output = arg 110 | elif opt in ("-p", "--patience"): 111 | patience = arg 112 | elif opt in ("-v", "--visualize"): 113 | visualize = True 114 | 115 | if '--net' not in sys.argv or '-n' not in sys.argv: 116 | from .models.UNetEx import UNetEx 117 | net = UNetEx 118 | 119 | options = { 120 | 'device': device, 121 | 'net': net, 122 | 'model_input': model_input, 123 | 'model_output': model_output, 124 | 'output': output, 125 | 'kernel_size': kernel_size, 126 | 'filters': filters, 127 | 'learning_rate': learning_rate, 128 | 'epochs': epochs, 129 | 'batch_size': batch_size, 130 | 'patience': patience, 131 | 'visualize': visualize, 132 | } 133 | 134 | return options 135 | 136 | def main(): 137 | options = parseOpts(sys.argv[1:]) 138 | 139 | x = pickle.load(open(options["model_input"], "rb")) 140 | y = pickle.load(open(options["model_output"], "rb")) 141 | 142 | # Shuffle the data 143 | indices = list(range(len(x))) 144 | random.shuffle(indices) 145 | x = x[indices] 146 | y = y[indices] 147 | 148 | x = torch.FloatTensor(x) 149 | y = torch.FloatTensor(y) 150 | 151 | batch = x.shape[0] 152 | nx = x.shape[2] 153 | ny = x.shape[3] 154 | 155 | channels_weights = torch.sqrt(torch.mean(y.permute(0, 2, 3, 1) 156 | .reshape((batch*nx*ny,3)) ** 2, dim=0)).view(1, -1, 1, 1).to(options["device"]) 157 | 158 | dirname = os.path.dirname(os.path.abspath(options["output"])) 159 | if dirname and not os.path.exists(dirname): 160 | os.makedirs(dirname, exist_ok=True) 161 | 162 | # Spliting dataset into 70% train and 30% test 163 | train_data, test_data = split_tensors(x, y, ratio=0.7) 164 | 165 | train_dataset, test_dataset = TensorDataset(*train_data), TensorDataset(*test_data) 166 | test_x, test_y = test_dataset[:] 167 | 168 | torch.manual_seed(0) 169 | 170 | model = options["net"]( 171 | 3, 172 | 3, 173 | filters=options["filters"], 174 | kernel_size=options["kernel_size"], 175 | batch_norm=False, 176 | weight_norm=False 177 | ) 178 | 179 | # Define optimizer 180 | optimizer = torch.optim.AdamW( 181 | model.parameters(), 182 | lr=options["learning_rate"], 183 | weight_decay=0.005 184 | ) 185 | 186 | config = {} 187 | train_loss_curve = [] 188 | test_loss_curve = [] 189 | train_mse_curve = [] 190 | test_mse_curve = [] 191 | train_ux_curve = [] 192 | test_ux_curve = [] 193 | train_uy_curve = [] 194 | test_uy_curve = [] 195 | train_p_curve = [] 196 | test_p_curve = [] 197 | 198 | def after_epoch(scope): 199 | train_loss_curve.append(scope["train_loss"]) 200 | test_loss_curve.append(scope["val_loss"]) 201 | train_mse_curve.append(scope["train_metrics"]["mse"]) 202 | test_mse_curve.append(scope["val_metrics"]["mse"]) 203 | train_ux_curve.append(scope["train_metrics"]["ux"]) 204 | test_ux_curve.append(scope["val_metrics"]["ux"]) 205 | train_uy_curve.append(scope["train_metrics"]["uy"]) 206 | test_uy_curve.append(scope["val_metrics"]["uy"]) 207 | train_p_curve.append(scope["train_metrics"]["p"]) 208 | test_p_curve.append(scope["val_metrics"]["p"]) 209 | 210 | def loss_func(model, batch): 211 | x, y = batch 212 | output = model(x) 213 | lossu = ((output[:,0,:,:] - y[:,0,:,:]) ** 2).reshape( 214 | (output.shape[0],1,output.shape[2],output.shape[3])) 215 | lossv = ((output[:,1,:,:] - y[:,1,:,:]) ** 2).reshape( 216 | (output.shape[0],1,output.shape[2],output.shape[3])) 217 | lossp = torch.abs((output[:,2,:,:] - y[:,2,:,:])).reshape( 218 | (output.shape[0],1,output.shape[2],output.shape[3])) 219 | loss = (lossu + lossv + lossp)/channels_weights 220 | 221 | return torch.sum(loss), output 222 | 223 | # Training model 224 | DeepCFD, train_metrics, train_loss, test_metrics, test_loss = train_model( 225 | model, 226 | loss_func, 227 | train_dataset, 228 | test_dataset, 229 | optimizer, 230 | epochs=options["epochs"], 231 | batch_size=options["batch_size"], 232 | device=options["device"], 233 | m_mse_name="Total MSE", 234 | m_mse_on_batch=lambda scope: 235 | float(torch.sum((scope["output"] - scope["batch"][1]) ** 2)), 236 | m_mse_on_epoch=lambda scope: 237 | sum(scope["list"]) / len(scope["dataset"]), 238 | m_ux_name="Ux MSE", 239 | m_ux_on_batch=lambda scope: 240 | float(torch.sum((scope["output"][:,0,:,:] - 241 | scope["batch"][1][:,0,:,:]) ** 2)), 242 | m_ux_on_epoch=lambda scope: 243 | sum(scope["list"]) / len(scope["dataset"]), 244 | m_uy_name="Uy MSE", 245 | m_uy_on_batch=lambda scope: 246 | float(torch.sum((scope["output"][:,1,:,:] - 247 | scope["batch"][1][:,1,:,:]) ** 2)), 248 | m_uy_on_epoch=lambda scope: 249 | sum(scope["list"]) / len(scope["dataset"]), 250 | m_p_name="p MSE", 251 | m_p_on_batch=lambda scope: 252 | float(torch.sum((scope["output"][:,2,:,:] - 253 | scope["batch"][1][:,2,:,:]) ** 2)), 254 | m_p_on_epoch=lambda scope: 255 | sum(scope["list"]) / 256 | len(scope["dataset"]), patience=options["patience"], after_epoch=after_epoch 257 | ) 258 | 259 | state_dict = DeepCFD.state_dict() 260 | state_dict["input_shape"] = (1, 3, nx, ny) 261 | state_dict["filters"] = options["filters"] 262 | state_dict["kernel_size"] = options["kernel_size"] 263 | state_dict["architecture"] = options["net"] 264 | 265 | torch.save(state_dict, options["output"]) 266 | 267 | if (options["visualize"]): 268 | out = DeepCFD(test_x[:10].to(options["device"])) 269 | error = torch.abs(out.cpu() - test_y[:10].cpu()) 270 | s = 0 271 | visualize( 272 | test_y[:10].cpu().detach().numpy(), 273 | out[:10].cpu().detach().numpy(), 274 | error[:10].cpu().detach().numpy(), 275 | s 276 | ) 277 | 278 | if __name__ == "__main__": 279 | main() 280 | -------------------------------------------------------------------------------- /src/deepcfd/functions.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import numpy as np 3 | from matplotlib import pyplot as plt 4 | 5 | def split_tensors(*tensors, ratio): 6 | assert len(tensors) > 0 7 | split1, split2 = [], [] 8 | count = len(tensors[0]) 9 | for tensor in tensors: 10 | assert len(tensor) == count 11 | split1.append(tensor[:int(len(tensor) * ratio)]) 12 | split2.append(tensor[int(len(tensor) * ratio):]) 13 | if len(tensors) == 1: 14 | split1, split2 = split1[0], split2[0] 15 | return split1, split2 16 | 17 | def initialize(model, gain=1, std=0.02): 18 | for module in model.modules(): 19 | if type(module) in [nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d]: 20 | nn.init.xavier_normal_(module.weight, gain) 21 | if module.bias is not None: 22 | nn.init.normal_(module.bias, 0, std) 23 | 24 | def visualize(sample_y, out_y, error, s): 25 | minu = np.min(sample_y[s, 0, :, :]) 26 | maxu = np.max(sample_y[s, 0, :, :]) 27 | 28 | minv = np.min(sample_y[s, 1, :, :]) 29 | maxv = np.max(sample_y[s, 1, :, :]) 30 | 31 | minp = np.min(sample_y[s, 2, :, :]) 32 | maxp = np.max(sample_y[s, 2, :, :]) 33 | 34 | mineu = np.min(error[s, 0, :, :]) 35 | maxeu = np.max(error[s, 0, :, :]) 36 | 37 | minev = np.min(error[s, 1, :, :]) 38 | maxev = np.max(error[s, 1, :, :]) 39 | 40 | minep = np.min(error[s, 2, :, :]) 41 | maxep = np.max(error[s, 2, :, :]) 42 | 43 | nx = sample_y.shape[2] 44 | ny = sample_y.shape[3] 45 | 46 | plot_options = {'cmap': 'jet', 'origin': 'lower', 'extent': [0,nx,0,ny]} 47 | 48 | plt.figure() 49 | fig = plt.gcf() 50 | fig.set_size_inches(15, 10) 51 | plt.subplot(3, 3, 1) 52 | plt.title('CFD', fontsize=18) 53 | plt.imshow(np.transpose(sample_y[s, 0, :, :]), vmin = minu, vmax = maxu, **plot_options) 54 | plt.colorbar(orientation='horizontal') 55 | plt.ylabel('Ux', fontsize=18) 56 | plt.subplot(3, 3, 2) 57 | plt.title('CNN', fontsize=18) 58 | plt.imshow(np.transpose(out_y[s, 0, :, :]), vmin = minu, vmax =maxu, **plot_options) 59 | plt.colorbar(orientation='horizontal') 60 | plt.subplot(3, 3, 3) 61 | plt.title('Error', fontsize=18) 62 | plt.imshow(np.transpose(error[s, 0, :, :]), vmin = mineu, vmax = maxeu, **plot_options) 63 | plt.colorbar(orientation='horizontal') 64 | 65 | plt.subplot(3, 3, 4) 66 | plt.imshow(np.transpose(sample_y[s, 1, :, :]), vmin = minv, vmax = maxv, **plot_options) 67 | plt.colorbar(orientation='horizontal') 68 | plt.ylabel('Uy', fontsize=18) 69 | plt.subplot(3, 3, 5) 70 | plt.imshow(np.transpose(out_y[s, 1, :, :]), vmin = minv, vmax = maxv, **plot_options) 71 | plt.colorbar(orientation='horizontal') 72 | plt.subplot(3, 3, 6) 73 | plt.imshow(np.transpose(error[s, 1, :, :]), vmin = minev, vmax = maxev, **plot_options) 74 | plt.colorbar(orientation='horizontal') 75 | 76 | plt.subplot(3, 3, 7) 77 | plt.imshow(np.transpose(sample_y[s, 2, :, :]), vmin = minp, vmax = maxp, **plot_options) 78 | plt.colorbar(orientation='horizontal') 79 | plt.ylabel('p', fontsize=18) 80 | plt.subplot(3, 3, 8) 81 | plt.imshow(np.transpose(out_y[s, 2, :, :]), vmin = minp, vmax = maxp, **plot_options) 82 | plt.colorbar(orientation='horizontal') 83 | plt.subplot(3, 3, 9) 84 | plt.imshow(np.transpose(error[s, 2, :, :]), vmin = minep, vmax = maxep, **plot_options) 85 | plt.colorbar(orientation='horizontal') 86 | plt.tight_layout() 87 | plt.show() 88 | -------------------------------------------------------------------------------- /src/deepcfd/models/AutoEncoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | 6 | 7 | def create_layer(in_channels, out_channels, kernel_size, wn=True, bn=True, 8 | activation=nn.ReLU, convolution=nn.Conv2d): 9 | assert kernel_size % 2 == 1 10 | layer = [] 11 | conv = convolution(in_channels, out_channels, kernel_size, padding=kernel_size // 2) 12 | if wn: 13 | conv = weight_norm(conv) 14 | layer.append(conv) 15 | if activation is not None: 16 | layer.append(activation()) 17 | if bn: 18 | layer.append(nn.BatchNorm2d(out_channels)) 19 | return nn.Sequential(*layer) 20 | 21 | 22 | class AutoEncoder(nn.Module): 23 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], 24 | weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): 25 | super().__init__() 26 | assert len(filters) > 0 27 | encoder = [] 28 | decoder = [] 29 | for i in range(len(filters)): 30 | if i == 0: 31 | encoder_layer = create_layer(in_channels, filters[i], kernel_size, weight_norm, batch_norm, activation, nn.Conv2d) 32 | decoder_layer = create_layer(filters[i], out_channels, kernel_size, weight_norm, False, final_activation, nn.ConvTranspose2d) 33 | else: 34 | encoder_layer = create_layer(filters[i-1], filters[i], kernel_size, weight_norm, batch_norm, activation, nn.Conv2d) 35 | decoder_layer = create_layer(filters[i], filters[i-1], kernel_size, weight_norm, batch_norm, activation, nn.ConvTranspose2d) 36 | encoder = encoder + [encoder_layer] 37 | decoder = [decoder_layer] + decoder 38 | self.encoder = nn.Sequential(*encoder) 39 | self.decoder = nn.Sequential(*decoder) 40 | 41 | def forward(self, x): 42 | return self.decoder(self.encoder(x)) 43 | -------------------------------------------------------------------------------- /src/deepcfd/models/AutoEncoderEx.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | from .AutoEncoder import create_layer 6 | 7 | 8 | class AutoEncoderEx(nn.Module): 9 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[8, 16, 16, 32], 10 | weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): 11 | super().__init__() 12 | assert len(filters) > 0 13 | encoder = [] 14 | decoder = [[] for _ in range(out_channels)] 15 | for i in range(len(filters)): 16 | if i == 0: 17 | encoder_layer = create_layer(in_channels, filters[i], kernel_size, weight_norm, batch_norm, activation, nn.Conv2d) 18 | decoder_layer = [create_layer(filters[i], 1, kernel_size, weight_norm, False, final_activation, nn.ConvTranspose2d) for _ in range(out_channels)] 19 | else: 20 | encoder_layer = create_layer(filters[i-1], filters[i], kernel_size, weight_norm, batch_norm, activation, nn.Conv2d) 21 | decoder_layer = [create_layer(filters[i], filters[i-1], kernel_size, weight_norm, batch_norm, activation, nn.ConvTranspose2d) for _ in range(out_channels)] 22 | encoder = encoder + [encoder_layer] 23 | for c in range(out_channels): 24 | decoder[c] = [decoder_layer[c]] + decoder[c] 25 | self.encoder = nn.Sequential(*encoder) 26 | for c in range(out_channels): 27 | decoder[c] = nn.Sequential(*decoder[c]) 28 | self.decoder = nn.Sequential(*decoder) 29 | 30 | def forward(self, x): 31 | x = self.encoder(x) 32 | y = [] 33 | for c in self.decoder: 34 | y.append(c(x)) 35 | return torch.cat(y, dim=1) 36 | -------------------------------------------------------------------------------- /src/deepcfd/models/UNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | from .AutoEncoder import create_layer 6 | 7 | 8 | def create_encoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 9 | activation=nn.ReLU, layers=2): 10 | encoder = [] 11 | for i in range(layers): 12 | _in = out_channels 13 | _out = out_channels 14 | if i == 0: 15 | _in = in_channels 16 | encoder.append(create_layer(_in, _out, kernel_size, wn, bn, activation, nn.Conv2d)) 17 | return nn.Sequential(*encoder) 18 | 19 | 20 | def create_decoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 21 | activation=nn.ReLU, layers=2, final_layer=False): 22 | decoder = [] 23 | for i in range(layers): 24 | _in = in_channels 25 | _out = in_channels 26 | _bn = bn 27 | _activation = activation 28 | if i == 0: 29 | _in = in_channels * 2 30 | if i == layers - 1: 31 | _out = out_channels 32 | if final_layer: 33 | _bn = False 34 | _activation = None 35 | decoder.append(create_layer(_in, _out, kernel_size, wn, _bn, _activation, nn.ConvTranspose2d)) 36 | return nn.Sequential(*decoder) 37 | 38 | 39 | def create_encoder(in_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 40 | encoder = [] 41 | for i in range(len(filters)): 42 | if i == 0: 43 | encoder_layer = create_encoder_block(in_channels, filters[i], kernel_size, wn, bn, activation, layers) 44 | else: 45 | encoder_layer = create_encoder_block(filters[i-1], filters[i], kernel_size, wn, bn, activation, layers) 46 | encoder = encoder + [encoder_layer] 47 | return nn.Sequential(*encoder) 48 | 49 | 50 | def create_decoder(out_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 51 | decoder = [] 52 | for i in range(len(filters)): 53 | if i == 0: 54 | decoder_layer = create_decoder_block(filters[i], out_channels, kernel_size, wn, bn, activation, layers, final_layer=True) 55 | else: 56 | decoder_layer = create_decoder_block(filters[i], filters[i-1], kernel_size, wn, bn, activation, layers, final_layer=False) 57 | decoder = [decoder_layer] + decoder 58 | return nn.Sequential(*decoder) 59 | 60 | 61 | class UNet(nn.Module): 62 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=2, 63 | weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): 64 | super().__init__() 65 | assert len(filters) > 0 66 | self.final_activation = final_activation 67 | self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) 68 | self.decoder = create_decoder(out_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) 69 | 70 | def encode(self, x): 71 | tensors = [] 72 | indices = [] 73 | sizes = [] 74 | for encoder in self.encoder: 75 | x = encoder(x) 76 | sizes.append(x.size()) 77 | tensors.append(x) 78 | x, ind = F.max_pool2d(x, 2, 2, return_indices=True) 79 | indices.append(ind) 80 | return x, tensors, indices, sizes 81 | 82 | def decode(self, x, tensors, indices, sizes): 83 | for decoder in self.decoder: 84 | tensor = tensors.pop() 85 | size = sizes.pop() 86 | ind = indices.pop() 87 | x = F.max_unpool2d(x, ind, 2, 2, output_size=size) 88 | x = torch.cat([tensor, x], dim=1) 89 | x = decoder(x) 90 | return x 91 | 92 | def forward(self, x): 93 | x, tensors, indices, sizes = self.encode(x) 94 | x = self.decode(x, tensors, indices, sizes) 95 | if self.final_activation is not None: 96 | x = self.final_activation(x) 97 | return x 98 | -------------------------------------------------------------------------------- /src/deepcfd/models/UNetEx.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | from .AutoEncoder import create_layer 6 | 7 | 8 | def create_encoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 9 | activation=nn.ReLU, layers=2): 10 | encoder = [] 11 | for i in range(layers): 12 | _in = out_channels 13 | _out = out_channels 14 | if i == 0: 15 | _in = in_channels 16 | encoder.append(create_layer(_in, _out, kernel_size, wn, bn, activation, nn.Conv2d)) 17 | return nn.Sequential(*encoder) 18 | 19 | 20 | def create_decoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 21 | activation=nn.ReLU, layers=2, final_layer=False): 22 | decoder = [] 23 | for i in range(layers): 24 | _in = in_channels 25 | _out = in_channels 26 | _bn = bn 27 | _activation = activation 28 | if i == 0: 29 | _in = in_channels * 2 30 | if i == layers - 1: 31 | _out = out_channels 32 | if final_layer: 33 | _bn = False 34 | _activation = None 35 | decoder.append(create_layer(_in, _out, kernel_size, wn, _bn, _activation, nn.ConvTranspose2d)) 36 | return nn.Sequential(*decoder) 37 | 38 | 39 | def create_encoder(in_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 40 | encoder = [] 41 | for i in range(len(filters)): 42 | if i == 0: 43 | encoder_layer = create_encoder_block(in_channels, filters[i], kernel_size, wn, bn, activation, layers) 44 | else: 45 | encoder_layer = create_encoder_block(filters[i-1], filters[i], kernel_size, wn, bn, activation, layers) 46 | encoder = encoder + [encoder_layer] 47 | return nn.Sequential(*encoder) 48 | 49 | 50 | def create_decoder(out_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 51 | decoder = [] 52 | for i in range(len(filters)): 53 | if i == 0: 54 | decoder_layer = create_decoder_block(filters[i], out_channels, kernel_size, wn, bn, activation, layers, final_layer=True) 55 | else: 56 | decoder_layer = create_decoder_block(filters[i], filters[i-1], kernel_size, wn, bn, activation, layers, final_layer=False) 57 | decoder = [decoder_layer] + decoder 58 | return nn.Sequential(*decoder) 59 | 60 | 61 | class UNetEx(nn.Module): 62 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=3, 63 | weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): 64 | super().__init__() 65 | assert len(filters) > 0 66 | self.final_activation = final_activation 67 | self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) 68 | decoders = [] 69 | for i in range(out_channels): 70 | decoders.append(create_decoder(1, filters, kernel_size, weight_norm, batch_norm, activation, layers)) 71 | self.decoders = nn.Sequential(*decoders) 72 | 73 | def encode(self, x): 74 | tensors = [] 75 | indices = [] 76 | sizes = [] 77 | for encoder in self.encoder: 78 | x = encoder(x) 79 | sizes.append(x.size()) 80 | tensors.append(x) 81 | x, ind = F.max_pool2d(x, 2, 2, return_indices=True) 82 | indices.append(ind) 83 | return x, tensors, indices, sizes 84 | 85 | def decode(self, _x, _tensors, _indices, _sizes): 86 | y = [] 87 | for _decoder in self.decoders: 88 | x = _x 89 | tensors = _tensors[:] 90 | indices = _indices[:] 91 | sizes = _sizes[:] 92 | for decoder in _decoder: 93 | tensor = tensors.pop() 94 | size = sizes.pop() 95 | ind = indices.pop() 96 | x = F.max_unpool2d(x, ind, 2, 2, output_size=size) 97 | x = torch.cat([tensor, x], dim=1) 98 | x = decoder(x) 99 | y.append(x) 100 | return torch.cat(y, dim=1) 101 | 102 | def forward(self, x): 103 | x, tensors, indices, sizes = self.encode(x) 104 | x = self.decode(x, tensors, indices, sizes) 105 | if self.final_activation is not None: 106 | x = self.final_activation(x) 107 | return x 108 | -------------------------------------------------------------------------------- /src/deepcfd/models/UNetExAvg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | from .AutoEncoder import create_layer 6 | 7 | 8 | def create_encoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 9 | activation=nn.ReLU, layers=2): 10 | encoder = [] 11 | for i in range(layers): 12 | _in = out_channels 13 | _out = out_channels 14 | if i == 0: 15 | _in = in_channels 16 | encoder.append(create_layer(_in, _out, kernel_size, wn, bn, activation, nn.Conv2d)) 17 | return nn.Sequential(*encoder) 18 | 19 | 20 | def create_decoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 21 | activation=nn.ReLU, layers=2, final_layer=False): 22 | decoder = [] 23 | for i in range(layers): 24 | _in = in_channels 25 | _out = in_channels 26 | _bn = bn 27 | _activation = activation 28 | if i == 0: 29 | _in = in_channels * 2 30 | if i == layers - 1: 31 | _out = out_channels 32 | if final_layer: 33 | _bn = False 34 | _activation = None 35 | decoder.append(create_layer(_in, _out, kernel_size, wn, _bn, _activation, nn.ConvTranspose2d)) 36 | return nn.Sequential(*decoder) 37 | 38 | 39 | def create_encoder(in_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 40 | encoder = [] 41 | for i in range(len(filters)): 42 | if i == 0: 43 | encoder_layer = create_encoder_block(in_channels, filters[i], kernel_size, wn, bn, activation, layers) 44 | else: 45 | encoder_layer = create_encoder_block(filters[i-1], filters[i], kernel_size, wn, bn, activation, layers) 46 | encoder = encoder + [encoder_layer] 47 | return nn.Sequential(*encoder) 48 | 49 | 50 | def create_decoder(out_channels, filters, kernel_size, wn=True, bn=True, activation=nn.ReLU, layers=2): 51 | decoder = [] 52 | for i in range(len(filters)): 53 | if i == 0: 54 | decoder_layer = create_decoder_block(filters[i], out_channels, kernel_size, wn, bn, activation, layers, final_layer=True) 55 | else: 56 | decoder_layer = create_decoder_block(filters[i], filters[i-1], kernel_size, wn, bn, activation, layers, final_layer=False) 57 | decoder = [decoder_layer] + decoder 58 | return nn.Sequential(*decoder) 59 | 60 | 61 | class UNetEx(nn.Module): 62 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=2, 63 | weight_norm=True, batch_norm=True, activation=nn.ReLU, final_activation=None): 64 | super().__init__() 65 | assert len(filters) > 0 66 | self.final_activation = final_activation 67 | self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) 68 | decoders = [] 69 | for i in range(out_channels): 70 | decoders.append(create_decoder(1, filters, kernel_size, weight_norm, batch_norm, activation, layers)) 71 | self.decoders = nn.Sequential(*decoders) 72 | 73 | def encode(self, x): 74 | tensors = [] 75 | sizes = [] 76 | for encoder in self.encoder: 77 | x = encoder(x) 78 | sizes.append(x.size()) 79 | tensors.append(x) 80 | x = F.avg_pool2d(x, 2, 2) 81 | return x, tensors, sizes 82 | 83 | def decode(self, _x, _tensors, _sizes): 84 | y = [] 85 | for _decoder in self.decoders: 86 | x = _x 87 | tensors = _tensors[:] 88 | sizes = _sizes[:] 89 | for decoder in _decoder: 90 | tensor = tensors.pop() 91 | size = sizes.pop() 92 | x = F.interpolate(x, size=size[-2:]) 93 | x = torch.cat([tensor, x], dim=1) 94 | x = decoder(x) 95 | y.append(x) 96 | return torch.cat(y, dim=1) 97 | 98 | def forward(self, x): 99 | x, tensors, sizes = self.encode(x) 100 | x = self.decode(x, tensors, sizes) 101 | if self.final_activation is not None: 102 | x = self.final_activation(x) 103 | return x 104 | -------------------------------------------------------------------------------- /src/deepcfd/models/UNetExMod.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.utils import weight_norm 5 | from .AutoEncoder import create_layer 6 | 7 | 8 | def create_encoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 9 | activation=nn.LeakyReLU, layers=2): 10 | encoder = [] 11 | for i in range(layers): 12 | _in = out_channels 13 | _out = out_channels 14 | if i == 0: 15 | _in = in_channels 16 | encoder.append(create_layer(_in, _out, kernel_size, wn, bn, activation, nn.Conv2d)) 17 | return nn.Sequential(*encoder) 18 | 19 | 20 | def create_decoder_block(in_channels, out_channels, kernel_size, wn=True, bn=True, 21 | activation=nn.LeakyReLU, layers=2, final_layer=False): 22 | decoder = [] 23 | for i in range(layers): 24 | _in = in_channels 25 | _out = in_channels 26 | _bn = bn 27 | _activation = activation 28 | if i == 0: 29 | _in = in_channels * 2 30 | if i == layers - 1: 31 | _out = out_channels 32 | if final_layer: 33 | _bn = False 34 | _activation = None 35 | decoder.append(create_layer(_in, _out, kernel_size, wn, _bn, _activation, nn.ConvTranspose2d)) 36 | return nn.Sequential(*decoder) 37 | 38 | 39 | def create_encoder(in_channels, filters, kernel_size, wn=True, bn=True, activation=nn.LeakyReLU, layers=2): 40 | encoder = [] 41 | for i in range(len(filters)): 42 | if i == 0: 43 | encoder_layer = create_encoder_block(in_channels, filters[i], kernel_size, wn, bn, activation, layers) 44 | else: 45 | encoder_layer = create_encoder_block(filters[i-1], filters[i], kernel_size, wn, bn, activation, layers) 46 | encoder = encoder + [encoder_layer] 47 | return nn.Sequential(*encoder) 48 | 49 | 50 | def create_decoder(out_channels, filters, kernel_size, wn=True, bn=True, activation=nn.LeakyReLU, layers=2): 51 | decoder = [] 52 | for i in range(len(filters)): 53 | if i == 0: 54 | decoder_layer = create_decoder_block(filters[i], out_channels, kernel_size, wn, bn, activation, layers, final_layer=True) 55 | else: 56 | decoder_layer = create_decoder_block(filters[i], filters[i-1], kernel_size, wn, bn, activation, layers, final_layer=False) 57 | decoder = [decoder_layer] + decoder 58 | return nn.Sequential(*decoder) 59 | 60 | 61 | class UNetExMod(nn.Module): 62 | def __init__(self, in_channels, out_channels, kernel_size=3, filters=[16, 32, 64], layers=3, 63 | weight_norm=True, batch_norm=True, activation=nn.LeakyReLU, final_activation=None): 64 | super().__init__() 65 | assert len(filters) > 0 66 | self.final_activation = final_activation 67 | self.encoder = create_encoder(in_channels, filters, kernel_size, weight_norm, batch_norm, activation, layers) 68 | decoders = [] 69 | for i in range(out_channels): 70 | decoders.append(create_decoder(1, filters, kernel_size, weight_norm, batch_norm, activation, layers)) 71 | self.decoders = nn.Sequential(*decoders) 72 | 73 | def encode(self, x): 74 | tensors = [] 75 | indices = [] 76 | sizes = [] 77 | for encoder in self.encoder: 78 | x = encoder(x) 79 | sizes.append(x.size()) 80 | tensors.append(x) 81 | x, ind = F.max_pool2d(x, 2, 2, return_indices=True) 82 | indices.append(ind) 83 | return x, tensors, indices, sizes 84 | 85 | def decode(self, _x, _tensors, _indices, _sizes): 86 | y = [] 87 | for _decoder in self.decoders: 88 | x = _x 89 | tensors = _tensors[:] 90 | indices = _indices[:] 91 | sizes = _sizes[:] 92 | for decoder in _decoder: 93 | tensor = tensors.pop() 94 | size = sizes.pop() 95 | ind = indices.pop() 96 | x = F.max_unpool2d(x, ind, 2, 2, output_size=size) 97 | x = torch.cat([tensor, x], dim=1) 98 | x = decoder(x) 99 | y.append(x) 100 | return torch.cat(y, dim=1) 101 | 102 | def forward(self, x): 103 | x, tensors, indices, sizes = self.encode(x) 104 | x = self.decode(x, tensors, indices, sizes) 105 | if self.final_activation is not None: 106 | x = self.final_activation(x) 107 | return x 108 | -------------------------------------------------------------------------------- /src/deepcfd/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdribeiro/DeepCFD/f9cc152fe1190895cdeb18b102c2640bc5587deb/src/deepcfd/models/__init__.py -------------------------------------------------------------------------------- /src/deepcfd/pytorchtools.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | class EarlyStopping: 5 | """Early stops the training if validation loss doesn't improve after a given patience.""" 6 | def __init__(self, patience=7, verbose=False, delta=0): 7 | """ 8 | Args: 9 | patience (int): How long to wait after last time validation loss improved. 10 | Default: 7 11 | verbose (bool): If True, prints a message for each validation loss improvement. 12 | Default: False 13 | delta (float): Minimum change in the monitored quantity to qualify as an improvement. 14 | Default: 0 15 | """ 16 | self.patience = patience 17 | self.verbose = verbose 18 | self.counter = 0 19 | self.best_score = None 20 | self.early_stop = False 21 | self.val_loss_min = np.Inf 22 | self.delta = delta 23 | 24 | def __call__(self, val_loss, model): 25 | 26 | score = -val_loss 27 | 28 | if self.best_score is None: 29 | self.best_score = score 30 | self.save_checkpoint(val_loss, model) 31 | elif score < self.best_score - self.delta: 32 | self.counter += 1 33 | print(f'EarlyStopping counter: {self.counter} out of {self.patience}') 34 | if self.counter >= self.patience: 35 | self.early_stop = True 36 | else: 37 | self.best_score = score 38 | self.save_checkpoint(val_loss, model) 39 | self.counter = 0 40 | 41 | def save_checkpoint(self, val_loss, model): 42 | '''Saves model when validation loss decrease.''' 43 | if self.verbose: 44 | print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...') 45 | torch.save(model.state_dict(), 'checkpoint.pt') 46 | self.val_loss_min = val_loss -------------------------------------------------------------------------------- /src/deepcfd/train_functions.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import torch 3 | from .pytorchtools import EarlyStopping 4 | 5 | 6 | def generate_metrics_list(metrics_def): 7 | list = {} 8 | for name in metrics_def.keys(): 9 | list[name] = [] 10 | return list 11 | 12 | 13 | def epoch(scope, loader, on_batch=None, training=False): 14 | model = scope["model"] 15 | optimizer = scope["optimizer"] 16 | loss_func = scope["loss_func"] 17 | metrics_def = scope["metrics_def"] 18 | scope = copy.copy(scope) 19 | scope["loader"] = loader 20 | 21 | metrics_list = generate_metrics_list(metrics_def) 22 | total_loss = 0 23 | if training: 24 | model.train() 25 | else: 26 | model.eval() 27 | for tensors in loader: 28 | if "process_batch" in scope and scope["process_batch"] is not None: 29 | tensors = scope["process_batch"](tensors) 30 | if "device" in scope and scope["device"] is not None: 31 | tensors = [tensor.to(scope["device"]) for tensor in tensors] 32 | loss, output = loss_func(model, tensors) 33 | if training: 34 | optimizer.zero_grad() 35 | loss.backward() 36 | optimizer.step() 37 | total_loss += loss.item() 38 | scope["batch"] = tensors 39 | scope["loss"] = loss 40 | scope["output"] = output 41 | scope["batch_metrics"] = {} 42 | for name, metric in metrics_def.items(): 43 | value = metric["on_batch"](scope) 44 | scope["batch_metrics"][name] = value 45 | metrics_list[name].append(value) 46 | if on_batch is not None: 47 | on_batch(scope) 48 | scope["metrics_list"] = metrics_list 49 | metrics = {} 50 | for name in metrics_def.keys(): 51 | scope["list"] = scope["metrics_list"][name] 52 | metrics[name] = metrics_def[name]["on_epoch"](scope) 53 | return total_loss, metrics 54 | 55 | 56 | def train(scope, train_dataset, val_dataset, patience=10, batch_size=256, print_function=print, eval_model=None, 57 | on_train_batch=None, on_val_batch=None, on_train_epoch=None, on_val_epoch=None, after_epoch=None): 58 | 59 | early_stopping = EarlyStopping(patience, verbose=True) 60 | 61 | epochs = scope["epochs"] 62 | model = scope["model"] 63 | metrics_def = scope["metrics_def"] 64 | scope = copy.copy(scope) 65 | 66 | scope["best_train_metric"] = None 67 | scope["best_train_loss"] = float("inf") 68 | scope["best_val_metrics"] = None 69 | scope["best_val_loss"] = float("inf") 70 | scope["best_model"] = None 71 | 72 | train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 73 | val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False) 74 | skips = 0 75 | for epoch_id in range(1, epochs + 1): 76 | scope["epoch"] = epoch_id 77 | print_function("Epoch #" + str(epoch_id), flush=True) 78 | # Training 79 | scope["dataset"] = train_dataset 80 | train_loss, train_metrics = epoch(scope, train_loader, on_train_batch, training=True) 81 | scope["train_loss"] = train_loss 82 | scope["train_metrics"] = train_metrics 83 | print_function("\tTrain Loss = " + str(train_loss), flush=True) 84 | for name in metrics_def.keys(): 85 | print_function("\tTrain " + metrics_def[name]["name"] + " = " + str(train_metrics[name]), flush=True) 86 | if on_train_epoch is not None: 87 | on_train_epoch(scope) 88 | del scope["dataset"] 89 | # Validation 90 | scope["dataset"] = val_dataset 91 | with torch.no_grad(): 92 | val_loss, val_metrics = epoch(scope, val_loader, on_val_batch, training=False) 93 | scope["val_loss"] = val_loss 94 | scope["val_metrics"] = val_metrics 95 | print_function("\tValidation Loss = " + str(val_loss), flush=True) 96 | for name in metrics_def.keys(): 97 | print_function("\tValidation " + metrics_def[name]["name"] + " = " + str(val_metrics[name]), flush=True) 98 | if on_val_epoch is not None: 99 | on_val_epoch(scope) 100 | del scope["dataset"] 101 | # Selection 102 | is_best = None 103 | if eval_model is not None: 104 | is_best = eval_model(scope) 105 | if is_best is None: 106 | is_best = val_loss < scope["best_val_loss"] 107 | if is_best: 108 | scope["best_train_metric"] = train_metrics 109 | scope["best_train_loss"] = train_loss 110 | scope["best_val_metrics"] = val_metrics 111 | scope["best_val_loss"] = val_loss 112 | scope["best_model"] = copy.deepcopy(model) 113 | print_function("Model saved!", flush=True) 114 | skips = 0 115 | else: 116 | skips += 1 117 | if after_epoch is not None: 118 | after_epoch(scope) 119 | early_stopping(val_loss, scope["best_model"]) 120 | if early_stopping.early_stop: 121 | print_function("Early stopping", flush=True) 122 | break 123 | 124 | return scope["best_model"], scope["best_train_metric"], scope["best_train_loss"],\ 125 | scope["best_val_metrics"], scope["best_val_loss"] 126 | 127 | 128 | def train_model(model, loss_func, train_dataset, val_dataset, optimizer, process_batch=None, eval_model=None, 129 | on_train_batch=None, on_val_batch=None, on_train_epoch=None, on_val_epoch=None, after_epoch=None, 130 | epochs=100, batch_size=256, patience=10, device=0, **kwargs): 131 | model = model.to(device) 132 | scope = {} 133 | scope["model"] = model 134 | scope["loss_func"] = loss_func 135 | scope["train_dataset"] = train_dataset 136 | scope["val_dataset"] = val_dataset 137 | scope["optimizer"] = optimizer 138 | scope["process_batch"] = process_batch 139 | scope["epochs"] = epochs 140 | scope["batch_size"] = batch_size 141 | scope["device"] = device 142 | metrics_def = {} 143 | names = [] 144 | for key in kwargs.keys(): 145 | parts = key.split("_") 146 | if len(parts) == 3 and parts[0] == "m": 147 | if parts[1] not in names: 148 | names.append(parts[1]) 149 | for name in names: 150 | if "m_" + name + "_name" in kwargs and "m_" + name + "_on_batch" in kwargs and "m_" + name + "_on_epoch" in kwargs: 151 | metrics_def[name] = { 152 | "name": kwargs["m_" + name + "_name"], 153 | "on_batch": kwargs["m_" + name + "_on_batch"], 154 | "on_epoch": kwargs["m_" + name + "_on_epoch"], 155 | } 156 | else: 157 | print("Warning: " + name + " metric is incomplete!") 158 | scope["metrics_def"] = metrics_def 159 | return train(scope, train_dataset, val_dataset, eval_model=eval_model, on_train_batch=on_train_batch, 160 | on_val_batch=on_val_batch, on_train_epoch=on_train_epoch, on_val_epoch=on_val_epoch, after_epoch=after_epoch, 161 | batch_size=batch_size, patience=patience) 162 | --------------------------------------------------------------------------------