symcad.core.ML.NeuralNet

  1#!/usr/bin/env python3
  2# Copyright (C) 2022, Will Hedgecock
  3#
  4# This program is free software: you can redistribute it and/or modify
  5# it under the terms of the GNU General Public License as published by
  6# the Free Software Foundation, either version 3 of the License, or
  7# (at your option) any later version.
  8#
  9# This program is distributed in the hope that it will be useful,
 10# but WITHOUT ANY WARRANTY; without even the implied warranty of
 11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 12# GNU General Public License for more details.
 13#
 14# You should have received a copy of the GNU General Public License
 15# along with this program.  If not, see <http://www.gnu.org/licenses/>.
 16
 17from constraint_prog.sympy_func import NeuralFunc
 18from .NeuralNetTrainer import NeuralNetTrainer
 19from typing import Dict, List, Optional, Tuple, TypeVar, Union
 20from ..CAD import CadGeneral
 21from pathlib import Path
 22import io, tarfile, torch
 23
 24SymPart = TypeVar('SymPart', bound='SymPart')
 25
 26class NeuralNet(object):
 27   """Private container that houses all neural networks for a given `SymPart`."""
 28
 29   # Public attributes ----------------------------------------------------------------------------
 30
 31   networks: Dict[str, torch.ScriptModule]
 32   """Dictionary of neural networks corresponding to each learned geometric property."""
 33
 34   sympy_networks: Dict[str, NeuralFunc]
 35   """Dictionary of sympy wrappers for the neural networks corresponding to each learned
 36   geometric property."""
 37
 38   param_order: List[str]
 39   """List containing the expected input order of geometric parameters to each neural network."""
 40
 41   param_transformations: Dict[str, Tuple[float, float, float, float]]
 42   """Dictionary of bounds, scalers, and biases for each geometric parameter in the neural net."""
 43
 44
 45   # Constructor ----------------------------------------------------------------------------------
 46
 47   def __init__(self, part_type_name: str,
 48                      net_storage_path: str,
 49                      part: Optional[SymPart] = None,
 50                      auto_train_network: Optional[bool] = False) -> None:
 51      """Initializes a container to house all neural networks for the geometric properties of
 52      the given `part`.
 53
 54      If `part` is not specified or is `None`, or if `auto_train_network` is set to `False`, a
 55      `RuntimeError` will be raised if the specified neural network does not already exist. If
 56      `auto_train_network` is `True` and no trained neural network exists, a new neural network
 57      will be trained to learn all available geometric properties for the specified `part`.
 58      """
 59
 60      # Initialize NeuralNet data structure
 61      super().__init__()
 62      self.networks = {}
 63      self.sympy_networks = {}
 64      self.param_order = []
 65      self.param_transformations = {}
 66
 67      # Ensure that the neural network exists and has been trained
 68      net_file_path = CadGeneral.CAD_BASE_PATH.joinpath(net_storage_path).absolute().resolve()
 69      if not net_file_path.exists():
 70         net_file_path = Path(net_storage_path).absolute().resolve()
 71      if not net_file_path.exists():
 72         net_converted_path = Path('converted').joinpath(Path(net_storage_path).stem + '.tar.xz')
 73         net_file_path = CadGeneral.CAD_BASE_PATH.joinpath(net_converted_path)
 74      if not net_file_path.exists():
 75         if auto_train_network and part is not None:
 76            net_file_path = \
 77               CadGeneral.CAD_BASE_PATH.joinpath(net_storage_path).absolute().resolve()
 78            trainer = NeuralNetTrainer(part, ['xlen', 'ylen', 'zlen', 'cg_x', 'cg_y', 'cg_z',
 79                                              'cb_x', 'cb_y', 'cb_z', 'material_volume',
 80                                              'displaced_volume', 'surface_area'])
 81            trainer.learn_parameters(32)
 82            trainer.save(str(net_file_path))
 83         else:
 84            raise RuntimeError('No trained neural network exists at "{}". '
 85                               'Set "auto_train_network" to True or manually train the '
 86                               'required network to continue'.format(net_storage_path))
 87
 88      # Load and parse all trained neural networks within the specified tarball
 89      with tarfile.open(net_file_path, 'r:xz') as zip_file:
 90         for filename in zip_file.getnames():
 91            data = zip_file.extractfile(filename)
 92            if filename == 'param_order.txt':
 93               self.param_order = data.read().decode('utf-8').split(';')
 94            elif filename == 'param_stats.txt':
 95               stats = data.read().decode('utf-8').split(';')
 96               for stat in stats:
 97                  param, val = stat.split(':')
 98                  self.param_transformations[param] = eval(val)
 99            else:
100               network_bytes = io.BytesIO(data.read())
101               network_name = '.'.join(filename.split('.')[:-1])
102               self.networks[network_name] = torch.jit.load(network_bytes)
103               self.networks[network_name].eval()
104      for network_name in self.networks.keys():
105         self.sympy_networks[network_name] = type(part_type_name + '_' + network_name,
106                                                  (NeuralFunc,),
107                                                  {'arity': len(self.param_order),
108                                                   'network': self.networks[network_name]})
109
110
111   # Public methods -------------------------------------------------------------------------------
112
113   def evaluate(self, property: str, **kwargs) -> Union[NeuralFunc, float]:
114      """Evaluates the specified parameters in `**kwargs` through the neural network
115      corresponding to the indicated geometric `property`.
116
117      Parameters
118      ----------
119      property : `str`
120         Geometric property for which the neural network should be evaluated. Available
121         options (if a trained network exists for the corresponding property) are:
122         - Lengths: `xlen`, `ylen`, `zlen`
123         - Centers of Gravity: `cg_x`, `cg_y`, `cg_z`
124         - Centers of Buoyancy: `cb_x`, `cb_y`, `cb_z`
125         - Volume and Area: `material_volume`, `displaced_volume`, `surface_area`
126
127      **kwargs : `Dict`
128         Set of named parameters that define the underlying geometry of this part.
129
130      Returns
131      -------
132      `Union[NeuralFunc, float]`
133         The requested property as evaluated by the underlying neural network.
134      """
135      if all([isinstance(val, float) or isinstance(val, int) for val in kwargs.values()]):
136         inputs = torch.empty(len(self.param_order))
137         for idx, param in enumerate(self.param_order):
138            transforms = self.param_transformations[param]
139            inputs[idx] = (kwargs.get(param) * transforms[2]) + transforms[3]
140         return self.networks[property](inputs).item()
141      else:
142         inputs = []
143         for param in self.param_order:
144            transforms = self.param_transformations[param]
145            inputs.append((kwargs.get(param) * transforms[2]) + transforms[3])
146         return self.sympy_networks[property](*inputs)