diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71f5e16..f09dd92 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,8 +10,33 @@ on: - main jobs: - # docker-image: - # uses: ./.github/workflows/docker.yml + docker-image: + uses: ./.github/workflows/docker.yml api-doc: - # needs: [docker-image] - uses: ./.github/workflows/sphinx.yml \ No newline at end of file + uses: ./.github/workflows/sphinx.yml + linux: + runs-on: ubuntu-latest + needs: [docker-image] + container: + image: ghcr.io/llnl/gplasdi/lasdi_env:latest + options: --user 1001 --privileged + volumes: + - /mnt:/mnt + steps: + - name: Cancel previous runs + uses: styfle/cancel-workflow-action@0.11.0 + with: + access_token: ${{ github.token }} + # - name: Set Swap Space + # uses: pierotofy/set-swap-space@master + # with: + # swap-size-gb: 10 + - name: Check out LaSDI + uses: actions/checkout@v3 + - name: Build LaSDI + run: | + cd ${GITHUB_WORKSPACE} + pip install . + - name: Test CNN2D + run: | + pytest -vrx tests/test_CNN2D.py \ No newline at end of file diff --git a/docker/Dockerfile b/docker/Dockerfile index bb49ef6..0a45e12 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -23,8 +23,8 @@ RUN sudo apt-get install -yq python3-dev RUN sudo apt-get install -yq python3-pip RUN sudo apt-get install python-is-python3 RUN sudo python -m pip install --upgrade pip -RUN sudo python -m pip install sphinx sphinx-autoapi sphinx_rtd_theme -#RUN sudo pip3 install numpy scipy argparse tables PyYAML h5py pybind11 pytest mpi4py merlin +# RUN sudo python -m pip install sphinx sphinx-autoapi sphinx_rtd_theme +RUN sudo pip3 install torch==2.0.1 numpy==1.23.0 scikit-learn==1.3 scipy==1.10 pyyaml==6.0 matplotlib==3.8.0 argparse==1.1 h5py pytest pytest-cov # RUN sudo apt-get clean -q diff --git a/docs/source/conf.py b/docs/source/conf.py index 29afdaa..111cc6a 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -25,6 +25,7 @@ ] autoapi_dirs = ['../../src'] +autoapi_python_class_content = 'both' napoleon_google_docstring = False napoleon_use_param = False diff --git a/examples/burgers1d.ipynb b/examples/burgers1d.ipynb index 1cff602..f65534e 100644 --- a/examples/burgers1d.ipynb +++ b/examples/burgers1d.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "d67dad03-9d76-4891-82ff-7e19d1369a24", "metadata": {}, "outputs": [], @@ -78,7 +78,7 @@ "outputs": [], "source": [ "# Specify the restart file you have.\n", - "filename = 'lasdi_10_01_2024_17_09.npy'\n", + "filename = 'restarts/burgers1d.restart.npy'\n", "\n", "import yaml\n", "from lasdi.workflow import initialize_trainer\n", @@ -111,7 +111,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "dcdac0c2", "metadata": {}, "outputs": [], @@ -179,6 +179,28 @@ "n_coef = restart_file['latent_dynamics']['ncoefs']" ] }, + { + "cell_type": "markdown", + "id": "03a96b35", + "metadata": {}, + "source": [ + "# Loss history" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e0163864", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(1)\n", + "plt.loglog(trainer.training_loss)\n", + "plt.loglog(trainer.ae_loss)\n", + "plt.loglog(trainer.ld_loss)\n", + "plt.loglog(trainer.coef_loss)" + ] + }, { "cell_type": "markdown", "id": "1262a0c3", @@ -306,7 +328,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11c629e4", + "id": "ba098b6d", "metadata": {}, "outputs": [], "source": [] diff --git a/src/lasdi/gplasdi.py b/src/lasdi/gplasdi.py index 3e156d4..8bbd903 100644 --- a/src/lasdi/gplasdi.py +++ b/src/lasdi/gplasdi.py @@ -145,6 +145,11 @@ def __init__(self, physics, autoencoder, latent_dynamics, param_space, config): self.X_train = torch.Tensor([]) self.X_test = torch.Tensor([]) + self.training_loss = [] + self.ae_loss = [] + self.ld_loss = [] + self.coef_loss = [] + return def train(self): @@ -163,11 +168,6 @@ def train(self): n_train = ps.n_train() ld = self.latent_dynamics - self.training_loss = [] - self.ae_loss = [] - self.ld_loss = [] - self.coef_loss = [] - ''' determine number of iterations. Perform n_iter iterations until overall iterations hit max_iter. @@ -291,4 +291,9 @@ def load(self, dict_): self.optimizer.load_state_dict(dict_['optimizer']) if (self.device != 'cpu'): optimizer_to(self.optimizer, self.device) + + self.training_loss = dict_['training_loss'] + self.ae_loss = dict_['ae_loss'] + self.ld_loss = dict_['ld_loss'] + self.coef_loss = dict_['coeff_loss'] return diff --git a/src/lasdi/latent_space.py b/src/lasdi/latent_space.py deleted file mode 100644 index fdf8fef..0000000 --- a/src/lasdi/latent_space.py +++ /dev/null @@ -1,175 +0,0 @@ -import torch -import numpy as np - -# activation dict -act_dict = {'ELU': torch.nn.ELU, - 'hardshrink': torch.nn.Hardshrink, - 'hardsigmoid': torch.nn.Hardsigmoid, - 'hardtanh': torch.nn.Hardtanh, - 'hardswish': torch.nn.Hardswish, - 'leakyReLU': torch.nn.LeakyReLU, - 'logsigmoid': torch.nn.LogSigmoid, - 'multihead': torch.nn.MultiheadAttention, - 'PReLU': torch.nn.PReLU, - 'ReLU': torch.nn.ReLU, - 'ReLU6': torch.nn.ReLU6, - 'RReLU': torch.nn.RReLU, - 'SELU': torch.nn.SELU, - 'CELU': torch.nn.CELU, - 'GELU': torch.nn.GELU, - 'sigmoid': torch.nn.Sigmoid, - 'SiLU': torch.nn.SiLU, - 'mish': torch.nn.Mish, - 'softplus': torch.nn.Softplus, - 'softshrink': torch.nn.Softshrink, - 'tanh': torch.nn.Tanh, - 'tanhshrink': torch.nn.Tanhshrink, - 'threshold': torch.nn.Threshold, - } - -def initial_condition_latent(param_grid, physics, autoencoder): - - ''' - - Outputs the initial condition in the latent space: Z0 = encoder(U0) - - ''' - - n_param = param_grid.shape[0] - Z0 = [] - - sol_shape = [1, 1] + physics.qgrid_size - - for i in range(n_param): - u0 = physics.initial_condition(param_grid[i]) - u0 = u0.reshape(sol_shape) - u0 = torch.Tensor(u0) - z0 = autoencoder.encoder(u0) - z0 = z0[0, 0, :].detach().numpy() - Z0.append(z0) - - return Z0 - -class MultiLayerPerceptron(torch.nn.Module): - - def __init__(self, layer_sizes, - act_type='sigmoid', reshape_index=None, reshape_shape=None, - threshold=0.1, value=0.0, num_heads=1): - super(MultiLayerPerceptron, self).__init__() - - # including input, hidden, output layers - self.n_layers = len(layer_sizes) - self.layer_sizes = layer_sizes - - # Linear features between layers - self.fcs = [] - for k in range(self.n_layers-1): - self.fcs += [torch.nn.Linear(layer_sizes[k], layer_sizes[k + 1])] - self.fcs = torch.nn.ModuleList(self.fcs) - self.init_weight() - - # Reshape input or output layer - assert((reshape_index is None) or (reshape_index in [0, -1])) - assert((reshape_shape is None) or (np.prod(reshape_shape) == layer_sizes[reshape_index])) - self.reshape_index = reshape_index - self.reshape_shape = reshape_shape - - # Initalize activation function - self.act_type = act_type - self.use_multihead = False - if act_type == "threshold": - self.act = act_dict[act_type](threshold, value) - - elif act_type == "multihead": - self.use_multihead = True - if (self.n_layers > 3): # if you have more than one hidden layer - self.act = [] - for i in range(self.n_layers-2): - self.act += [act_dict[act_type](layer_sizes[i+1], num_heads)] - else: - self.act = [torch.nn.Identity()] # No additional activation - self.act = torch.nn.ModuleList(self.fcs) - - #all other activation functions initialized here - else: - self.act = act_dict[act_type]() - return - - def forward(self, x): - if (self.reshape_index == 0): - # make sure the input has a proper shape - assert(list(x.shape[-len(self.reshape_shape):]) == self.reshape_shape) - # we use torch.Tensor.view instead of torch.Tensor.reshape, - # in order to avoid data copying. - x = x.view(list(x.shape[:-len(self.reshape_shape)]) + [self.layer_sizes[self.reshape_index]]) - - for i in range(self.n_layers-2): - x = self.fcs[i](x) # apply linear layer - if (self.use_multihead): - x = self.apply_attention(self, x, i) - else: - x = self.act(x) - - x = self.fcs[-1](x) - - if (self.reshape_index == -1): - # we use torch.Tensor.view instead of torch.Tensor.reshape, - # in order to avoid data copying. - x = x.view(list(x.shape[:-1]) + self.reshape_shape) - - return x - - def apply_attention(self, x, act_idx): - x = x.unsqueeze(1) # Add sequence dimension for attention - x, _ = self.act[act_idx](x, x, x) # apply attention - x = x.squeeze(1) # Remove sequence dimension - return x - - def init_weight(self): - # TODO(kevin): support other initializations? - for fc in self.fcs: - torch.nn.init.xavier_uniform_(fc.weight) - return - -class Autoencoder(torch.nn.Module): - - def __init__(self, physics, config): - super(Autoencoder, self).__init__() - - self.qgrid_size = physics.qgrid_size - self.space_dim = np.prod(self.qgrid_size) - hidden_units = config['hidden_units'] - n_z = config['latent_dimension'] - self.n_z = n_z - - layer_sizes = [self.space_dim] + hidden_units + [n_z] - #grab relevant initialization values from config - act_type = config['activation'] if 'activation' in config else 'sigmoid' - threshold = config["threshold"] if "threshold" in config else 0.1 - value = config["value"] if "value" in config else 0.0 - num_heads = config['num_heads'] if 'num_heads' in config else 1 - - self.encoder = MultiLayerPerceptron(layer_sizes, act_type, - reshape_index=0, reshape_shape=self.qgrid_size, - threshold=threshold, value=value, num_heads=num_heads) - - self.decoder = MultiLayerPerceptron(layer_sizes[::-1], act_type, - reshape_index=-1, reshape_shape=self.qgrid_size, - threshold=threshold, value=value, num_heads=num_heads) - - return - - def forward(self, x): - - x = self.encoder(x) - x = self.decoder(x) - - return x - - def export(self): - dict_ = {'autoencoder_param': self.cpu().state_dict()} - return dict_ - - def load(self, dict_): - self.load_state_dict(dict_['autoencoder_param']) - return \ No newline at end of file diff --git a/src/lasdi/latent_space/__init__.py b/src/lasdi/latent_space/__init__.py new file mode 100644 index 0000000..191d54b --- /dev/null +++ b/src/lasdi/latent_space/__init__.py @@ -0,0 +1,176 @@ +import torch +import numpy as np +from ..networks import MultiLayerPerceptron, CNN2D + +def initial_condition_latent(param_grid, physics, autoencoder): + + ''' + + Outputs the initial condition in the latent space: Z0 = encoder(U0) + + ''' + + n_param = param_grid.shape[0] + Z0 = [] + + sol_shape = [1, 1] + physics.qgrid_size + + for i in range(n_param): + u0 = physics.initial_condition(param_grid[i]) + u0 = u0.reshape(sol_shape) + u0 = torch.Tensor(u0) + z0 = autoencoder.encoder(u0) + z0 = z0[0, 0, :].detach().numpy() + Z0.append(z0) + + return Z0 + +class LatentSpace(torch.nn.Module): + + def __init__(self, physics, config): + super(LatentSpace, self).__init__() + + self.qgrid_size = physics.qgrid_size + self.n_z = config['latent_dimension'] + + return + + def forward(self, x): + raise RuntimeError("LatentSpace.forward: abstract method!") + + def export(self): + dict_ = {'qgrid_size': self.qgrid_size, + 'n_z': self.n_z} + return dict_ + + def load(self, dict_): + """ + Notes + ----- + This abstract class only checks if the variables in restart file are the same as the instance attributes. + """ + + assert(dict_['qgrid_size'] == self.qgrid_size) + assert(dict_['n_z'] == self.n_z) + return + + +class Autoencoder(LatentSpace): + + def __init__(self, physics, config): + super().__init__(physics, config) + + self.space_dim = np.prod(self.qgrid_size) + hidden_units = config['hidden_units'] + + layer_sizes = [self.space_dim] + hidden_units + [self.n_z] + #grab relevant initialization values from config + act_type = config['activation'] if 'activation' in config else 'sigmoid' + threshold = config["threshold"] if "threshold" in config else 0.1 + value = config["value"] if "value" in config else 0.0 + + self.encoder = MultiLayerPerceptron(layer_sizes, act_type, + reshape_index=0, reshape_shape=self.qgrid_size, + threshold=threshold, value=value) + + self.decoder = MultiLayerPerceptron(layer_sizes[::-1], act_type, + reshape_index=-1, reshape_shape=self.qgrid_size, + threshold=threshold, value=value) + + return + + def forward(self, x): + + x = self.encoder(x) + x = self.decoder(x) + + return x + + def export(self): + dict_ = super().export() + dict_['autoencoder_param'] = self.cpu().state_dict() + return dict_ + + def load(self, dict_): + super().load(dict_) + self.load_state_dict(dict_['autoencoder_param']) + return + +class Conv2DAutoencoder(LatentSpace): + def __init__(self, physics, config): + super().__init__(physics, config) + from lasdi.inputs import InputParser + parser = InputParser(config) + + assert(physics.dim == 2) + + if (len(self.qgrid_size) == 2): + cnn_layers = [[1] + self.qgrid_size] + cnn_layers += parser.getInput(['cnn_layers'], datatype=list) + + strides = parser.getInput(['strides'], fallback=[1] * (len(cnn_layers) - 1)) + paddings = parser.getInput(['paddings'], fallback=[0] * (len(cnn_layers) - 1)) + dilations = parser.getInput(['dilations'], fallback=[1] * (len(cnn_layers) - 1)) + + cnn_act_type = parser.getInput(['cnn_activation'], fallback='ReLU') + + batch_shape = parser.getInput(['batch_shape'], datatype=list) + data_shape = batch_shape + self.qgrid_size + + cnn_f = CNN2D(cnn_layers, 'forward', strides, paddings, + dilations, act_type=cnn_act_type, data_shape=data_shape) + cnn_b = CNN2D(cnn_layers[::-1], 'backward', strides[::-1], paddings[::-1], + dilations[::-1], act_type=cnn_act_type, data_shape=data_shape) + + mlp_layers = [np.prod(cnn_f.layer_sizes[-1])] + mlp_layers += parser.getInput(['mlp_layers'], datatype=list) + mlp_layers += [self.n_z] + + act_type = parser.getInput(['mlp_activation'], fallback='sigmoid') + threshold = parser.getInput(['threshold'], fallback=0.1) + value = parser.getInput(['value'], fallback=0.0) + + mlp_f = MultiLayerPerceptron(mlp_layers, act_type=act_type, + reshape_index=0, reshape_shape=cnn_f.layer_sizes[-1], + threshold=threshold, value=value) + mlp_b = MultiLayerPerceptron(mlp_layers[::-1], act_type=act_type, + reshape_index=-1, reshape_shape=cnn_b.layer_sizes[0], + threshold=threshold, value=value), + + self.encoder = torch.nn.Sequential(cnn_f, mlp_f) + self.decoder = torch.nn.Sequential(mlp_b, cnn_b) + + self.print_architecture() + + return + + def forward(self, x): + + x = self.encoder(x) + x = self.decoder(x) + + return x + + def export(self): + dict_ = {'autoencoder_param': self.cpu().state_dict()} + return dict_ + + def load(self, dict_): + self.load_state_dict(dict_['autoencoder_param']) + return + + def set_batch_shape(self, batch_shape): + data_shape = batch_shape + self.qgrid_size + + self.encoder[0].set_data_shape(data_shape) + self.decoder[1].set_data_shape(data_shape) + + self.print_architecture() + return + + def print_architecture(self): + self.encoder[0].print_data_shape() + self.encoder[1].print_architecture() + self.decoder[0].print_architecture() + self.decoder[1].print_data_shape() + return \ No newline at end of file diff --git a/src/lasdi/networks.py b/src/lasdi/networks.py new file mode 100644 index 0000000..f15195f --- /dev/null +++ b/src/lasdi/networks.py @@ -0,0 +1,648 @@ +import torch +import numpy as np + +act_dict = {'ELU': torch.nn.ELU, + 'hardshrink': torch.nn.Hardshrink, + 'hardsigmoid': torch.nn.Hardsigmoid, + 'hardtanh': torch.nn.Hardtanh, + 'hardswish': torch.nn.Hardswish, + 'leakyReLU': torch.nn.LeakyReLU, + 'logsigmoid': torch.nn.LogSigmoid, + 'multihead': torch.nn.MultiheadAttention, + 'PReLU': torch.nn.PReLU, + 'ReLU': torch.nn.ReLU, + 'ReLU6': torch.nn.ReLU6, + 'RReLU': torch.nn.RReLU, + 'SELU': torch.nn.SELU, + 'CELU': torch.nn.CELU, + 'GELU': torch.nn.GELU, + 'sigmoid': torch.nn.Sigmoid, + 'SiLU': torch.nn.SiLU, + 'mish': torch.nn.Mish, + 'softplus': torch.nn.Softplus, + 'softshrink': torch.nn.Softshrink, + 'tanh': torch.nn.Tanh, + 'tanhshrink': torch.nn.Tanhshrink, + 'threshold': torch.nn.Threshold, + } +""":obj:`dict` : Dictionary to activation functions. + +- :obj:`'ELU'`: :obj:`torch.nn.ELU` +- :obj:`'hardshrink'`: :obj:`torch.nn.Hardshrink` +- :obj:`'hardsigmoid'`: :obj:`torch.nn.Hardsigmoid` +- :obj:`'hardtanh'`: :obj:`torch.nn.Hardtanh` +- :obj:`'hardswish'`: :obj:`torch.nn.Hardswish` +- :obj:`'leakyReLU'`: :obj:`torch.nn.LeakyReLU` +- :obj:`'logsigmoid'`: :obj:`torch.nn.LogSigmoid` +- :obj:`'multihead'`: :obj:`torch.nn.MultiheadAttention` +- :obj:`'PReLU'`: :obj:`torch.nn.PReLU` +- :obj:`'ReLU'`: :obj:`torch.nn.ReLU` +- :obj:`'ReLU6'`: :obj:`torch.nn.ReLU6` +- :obj:`'RReLU'`: :obj:`torch.nn.RReLU` +- :obj:`'SELU'`: :obj:`torch.nn.SELU` +- :obj:`'CELU'`: :obj:`torch.nn.CELU` +- :obj:`'GELU'`: :obj:`torch.nn.GELU` +- :obj:`'sigmoid'`: :obj:`torch.nn.Sigmoid` +- :obj:`'SiLU'`: :obj:`torch.nn.SiLU` +- :obj:`'mish'`: :obj:`torch.nn.Mish` +- :obj:`'softplus'`: :obj:`torch.nn.Softplus` +- :obj:`'softshrink'`: :obj:`torch.nn.Softshrink` +- :obj:`'tanh'`: :obj:`torch.nn.Tanh` +- :obj:`'tanhshrink'`: :obj:`torch.nn.Tanhshrink` +- :obj:`'threshold'`: :obj:`torch.nn.Threshold` +""" + +class MultiLayerPerceptron(torch.nn.Module): + """Vanilla multi-layer perceptron neural networks module. + """ + + def __init__(self, layer_sizes, + act_type='sigmoid', reshape_index=None, reshape_shape=None, + threshold=0.1, value=0.0): + """ + Parameters + ---------- + layer_sizes : :obj:`list(int)` + List of vector dimensions of layers. + act_type : :obj:`str`, optional + Type of activation functions. By default :obj:`'sigmoid'` is used. + See :obj:`act_dict` for available types. + reshape_index : :obj:`int`, optinal + Index of layer to reshape input/output data. Either 0 or -1 is allowed. + + - 0 : the first (input) layer + - -1 : the last (output) layer + + By default the index is :obj:`None`, and reshaping is not executed. + reshape_shape : :obj:`list(int)`, optional + Target shape from/to which input/output data is reshaped. + Reshaping behavior changes by :attr:`reshape_index`. + By default the index is :obj:`None`, and reshaping is not executed. + For details on reshaping action, see :attr:`reshape_shape`. + + Note + ---- + :obj:`numpy.prod(reshape_shape) == layer_sizes[reshape_index]` + + """ + super(MultiLayerPerceptron, self).__init__() + + self.n_layers = len(layer_sizes) + """:obj:`int` : Depth of layers including input, hidden, output layers.""" + + self.layer_sizes = layer_sizes + """:obj:`list(int)` : Vector dimensions corresponding to each layer.""" + + self.fcs = [] + """:obj:`torch.nn.ModuleList` : linear features between layers.""" + for k in range(self.n_layers-1): + self.fcs += [torch.nn.Linear(layer_sizes[k], layer_sizes[k + 1])] + self.fcs = torch.nn.ModuleList(self.fcs) + + self.init_weight() + + # Reshape input or output layer + assert(reshape_index in [0, -1, None]) + assert((reshape_shape is None) or (np.prod(reshape_shape) == layer_sizes[reshape_index])) + self.reshape_index = reshape_index + """:obj:`int` : Index of layer to reshape input/output data. + + - 0 : the first (input) layer + - -1 : the last (output) layer + - :obj:`None` : no reshaping + """ + self.reshape_shape = reshape_shape + """:obj:`list(int)` : Target shape from/to which input/output data is reshaped. + For a reshape_shape :math:`[R_1, R_2, \ldots, R_n]`, + + - if :attr:`reshape_index` is 0, the input data shape is changed as + + .. math:: + [\ldots, R_1, R_2, \ldots, R_n] \\longrightarrow [\ldots, \prod_{i=1}^n R_i] + + - if :attr:`reshape_index` is -1, the output data shape is changed as + + .. math:: + [\ldots, \prod_{i=1}^n R_i] \\longrightarrow [\ldots, R_1, R_2, \ldots, R_n] + + - :obj:`None` : no reshaping + """ + + # Initalize activation function + self.act_type = act_type + """:obj:`str` : Type of activation functions.""" + + if act_type == "threshold": + act = act_dict[act_type](threshold, value) + """:obj:`torch.nn.Module` : Activation function used throughout the layers.""" + + elif act_type == "multihead": + raise RuntimeError("MultiLayerPerceptron: MultiheadAttention requires a different architecture!") + + #all other activation functions initialized here + else: + act = act_dict[act_type]() + + self.act = act + """:obj:`torch.nn.Module` : Activation function used throughout the layers.""" + return + + def forward(self, x): + """Evaluate through the module. + + Parameters + ---------- + x : :obj:`torch.Tensor` + Input data to pass into the module. + + Note + ---- + For :attr:`reshape_index` =0, + the last :math:`n` dimensions of :obj:`x` must match + :attr:`reshape_shape` :math:`=[R_1, R_2, \ldots, R_n]`. + + Returns + ------- + :obj:`torch.Tensor` + Output tensor evaluated from the module. + + Note + ---- + For :attr:`reshape_index` =-1, + the last dimension of the output tensor will be reshaped as + :attr:`reshape_shape` :math:`=[R_1, R_2, \ldots, R_n]`. + + """ + if (self.reshape_index == 0): + # make sure the input has a proper shape + assert(list(x.shape[-len(self.reshape_shape):]) == self.reshape_shape) + # we use torch.Tensor.view instead of torch.Tensor.reshape, + # in order to avoid data copying. + x = x.view(list(x.shape[:-len(self.reshape_shape)]) + [self.layer_sizes[self.reshape_index]]) + + for i in range(self.n_layers-2): + x = self.fcs[i](x) # apply linear layer + x = self.act(x) + + x = self.fcs[-1](x) + + if (self.reshape_index == -1): + # we use torch.Tensor.view instead of torch.Tensor.reshape, + # in order to avoid data copying. + x = x.view(list(x.shape[:-1]) + self.reshape_shape) + + return x + + def init_weight(self): + """Initialize weights of linear features according to Xavier uniform distribution.""" + + # TODO(kevin): support other initializations? + for fc in self.fcs: + torch.nn.init.xavier_uniform_(fc.weight) + return + + def print_architecture(self): + """Print out the architecture of the module.""" + print(self.layer_sizes) + +class CNN2D(torch.nn.Module): + """Two-dimensional convolutional neural networks.""" + + from enum import Enum + class Mode(Enum): + """Enumeration to specify direction of CNN.""" + Forward = 1 + """Contracting direction""" + Backward = -1 + """Expanding direction""" + + def __init__(self, layer_sizes, mode, + strides, paddings, dilations, + groups=1, bias=True, padding_mode='zeros', + act_type='ReLU', data_shape=None): + """ + Parameters + ---------- + layer_sizes : :obj:`numpy.array` + 2d array of tensor dimension of each layer. + See :attr:`layer_sizes`. + mode : :obj:`str` + Direction of CNN + - `forward`: contracting direction + - `backward`: expanding direction + strides : :obj:`list` + List of strides corresponding to each layer. + Each stride is either integer or tuple. + paddings : :obj:`list` + List of paddings corresponding to each layer. + Each padding is either integer or tuple. + dilations : :obj:`list` + List of dilations corresponding to each layer. + Each dilation is either integer or tuple. + groups : :obj:`int`, optional + Groups that applies to all layers. By default 1 + bias : :obj:`bool`, optional + Bias that applies to all layers. By default :obj:`True` + padding_mode : :obj:`str`, optional + Padding_mode that applies to all layers. By default :obj:`'zeros'` + act_type : :obj:`str`, optional + Activation function applied between all layers. By default :obj:`'ReLU'`. + See :obj:`act_dict` for available types. + data_shape : :obj:`list(int)`, optional + Data shape to/from which output/input data is reshaped. + See :attr:`data_shape` for details. + + Note + ---- + :obj:`len(strides) == layer_sizes.shape[0] - 1` + + :obj:`len(paddings) == layer_sizes.shape[0] - 1` + + :obj:`len(dilations) == layer_sizes.shape[0] - 1` + """ + super(CNN2D, self).__init__() + + if (mode == 'forward'): + self.mode = self.Mode.Forward + module = torch.nn.Conv2d + elif (mode == 'backward'): + self.mode = self.Mode.Backward + module = torch.nn.ConvTranspose2d + else: + raise RuntimeError('CNN2D: Unknown mode %s!' % mode) + + self.n_layers = len(layer_sizes) + """:obj:`int` : Depth of layers including input, hidden, output layers.""" + + self.layer_sizes = layer_sizes + """:obj:`numpy.array` : 2d integer array of shape :math:`[n\_layers, 3]`, + indicating tensor dimension of each layer. + For :math:`k`-th layer, the tensor dimension is + + .. math:: + layer\_sizes[k] = [channels, height, width] + + """ + + self.channels = [layer_sizes[k][0] for k in range(self.n_layers)] + """:obj:`list(int)` : list of channel size that + determines architecture of each layer. + For details on how architecture is determined, + see `torch API documentation `_. + """ + + # assert(len(kernel_sizes) == self.n_layers - 1) + assert(len(strides) == self.n_layers - 1) + assert(len(paddings) == self.n_layers - 1) + assert(len(dilations) == self.n_layers - 1) + # self.kernel_sizes = kernel_sizes + self.strides = strides + """:obj:`list` : list of strides that + determine architecture of each layer. + Each stride can be either integer or tuple. + For details on how architecture is determined, + see `torch API documentation `_. + """ + self.paddings = paddings + """:obj:`list` : list of paddings that + determine architecture of each layer. + Each padding can be either integer or tuple. + For details on how architecture is determined, + see `torch API documentation `_. + """ + self.dilations = dilations + """:obj:`list` : list of dilations that + determine architecture of each layer. + Each dilation can be either integer or tuple. + For details on how architecture is determined, + see `torch API documentation `_. + """ + + self.groups = groups + """:obj:`int` : groups that determine architecture of all layers. + For details on how architecture is determined, + see `torch API documentation `_. + """ + self.bias = bias + """:obj:`bool` : bias that determine architecture of all layers. + For details on how architecture is determined, + see `torch API documentation `_. + """ + self.padding_mode = padding_mode + """:obj:`str` : padding mode that determine architecture of all layers. + For details on how architecture is determined, + see `torch API documentation `_. + """ + + from lasdi.networks import act_dict + # TODO(kevin): not use threshold activation for now. + assert(act_type != 'threshold') + self.act = act_dict[act_type]() + """:obj:`torch.nn.Module` : activation function applied between all layers.""" + + self.kernel_sizes = [] + """:obj:`list` : list of kernel_sizes that + determine architecture of each layer. + Each kernel_size can be either integer or tuple. + Kernel size is automatically determined so that + output of the corresponding layer has the shape of the next layer. + + For details on how architecture is determined, + see `torch API documentation `_. + """ + self.fcs = [] + """:obj:`torch.nn.ModuleList` : module list of + :obj:`torch.nn.Conv2d` (forward) or :obj:`torch.nn.Conv2d` (backward).""" + + for k in range(self.n_layers - 1): + kernel_size = self.compute_kernel_size(self.layer_sizes[k][1:], self.layer_sizes[k+1][1:], + self.strides[k], self.paddings[k], self.dilations[k], self.mode) + out_shape = self.compute_output_layer_size(self.layer_sizes[k][1:], kernel_size, self.strides[k], + self.paddings[k], self.dilations[k], self.mode) + assert(self.layer_sizes[k+1][1:] == out_shape) + + self.kernel_sizes += [kernel_size] + self.fcs += [module(self.channels[k], self.channels[k+1], self.kernel_sizes[k], + stride=self.strides[k], padding=self.paddings[k], dilation=self.dilations[k], + groups=self.groups, bias=self.bias, padding_mode=self.padding_mode)] + + self.fcs = torch.nn.ModuleList(self.fcs) + self.init_weight() + + self.data_shape = data_shape + """:obj:`list(int)` : tensor dimension of the training data + that will be passed into/out of the module.""" + + self.batch_reshape = None + """:obj:`list(int)` : tensor dimension to which input/output data is reshaped. + + - Forward :attr:`mode`: shape of 3d-/4d-array + - Backward :attr:`mode`: shape of arbitrary nd-array + + Determined by :meth:`set_data_shape`. + """ + if (data_shape is not None): + self.set_data_shape(data_shape) + + return + + def set_data_shape(self, data_shape : list): + """ + Set the batch reshape in order to reshape the input/output batches + based on given training data shape. + + Forward :attr:`mode`: + + For :obj:`data_shape` :math:`=[N_1,\ldots,N_m]` + and the first layer size of :math:`[C_1, H_1, W_1]`, + + .. math:: + batch\_reshape = [R_1, C_1, H_1, W_1], + + where :math:`\prod_{i=1}^m N_i = R_1\\times C_1\\times H_1\\times W_1`. + + If :math:`m=2` and :math:`C_1=1`, then + + .. math:: + batch\_reshape = [C_1, H_1, W_1]. + + Note + ---- + For forward mode, :obj:`data_shape[-2:]==self.layer_sizes[0, 1:]` must be true. + + + Backward :attr:`mode`: + + :attr:`batch_shape` is the same as :obj:`data_shape`. + Output tensor of the module is reshaped as :obj:`data_shape`. + + Parameters + ---------- + data_shape : :obj:`list(int)` + Shape of the input/output data tensor for forward/backward mode. + """ + idx = 0 if (self.mode == CNN2D.Mode.Forward) else -1 + + if (self.mode == CNN2D.Mode.Forward): + if (len(data_shape) > 2): + if (data_shape[-3] != self.channels[idx]): + assert(self.channels[idx] == 1) + self.batch_reshape = [np.prod(data_shape[:-2]), 1] + else: + self.batch_reshape = [np.prod(data_shape[:-3]), self.channels[idx]] + elif (len(data_shape) == 2): + assert(self.channels[idx] == 1) + self.batch_reshape = [1] + + self.batch_reshape += data_shape[-2:] + else: + self.batch_reshape = list(data_shape) + + self.data_shape = list(data_shape) + return + + def print_data_shape(self): + """Print out the data shape and architecture of the module.""" + + mode_str = "forward" if (self.mode == CNN2D.Mode.Forward) else "backward" + print("mode: ", mode_str) + print("batch reshape: ", self.batch_reshape) + for k in range(self.n_layers - 1): + print('input layer: ', self.layer_sizes[k], + 'kernel size: ', self.kernel_sizes[k], + 'output layer: ', self.layer_sizes[k+1]) + return + + def forward(self, x): + """Evaluate through the module. + + Parameters + ---------- + x : :obj:`torch.nn.Tensor` + Input tensor to pass into the module. + + - Forward mode: nd array of shape :attr:`data_shape` + - Backward mode: Same shape as the output tensor of forward mode + + Returns + ------- + :obj:`torch.nn.Tensor` + Output tensor evaluated from the module. + + - Forward mode: 3d array of shape :obj:`self.layer_sizes[-1]`, + or 4d array of shape :obj:`[self.batch_reshape[0]] + self.layer_sizes[-1]` + - Backward mode: nd array of shape :attr:`data_shape` (equal to :attr:`batch_shape`) + """ + if ((self.batch_reshape is not None) and (self.mode == CNN2D.Mode.Forward)): + x = x.view(self.batch_reshape) + + for i in range(self.n_layers-2): + x = self.fcs[i](x) + x = self.act(x) + + x = self.fcs[-1](x) + + if ((self.batch_reshape is not None) and (self.mode == CNN2D.Mode.Backward)): + x = x.view(self.batch_reshape) + return x + + @classmethod + def compute_kernel_size(cls, input_shape, output_shape, stride, padding, dilation, mode): + """Compute kernel size that produces desired output shape from given input shape. + + The formula is based on torch API documentation + for `Conv2d `_ + and `ConvTranspose2d `_. + + Parameters + ---------- + input_shape : :obj:`int` or :obj:`tuple(int)` + output_shape : :obj:`int` or :obj:`tuple(int)` + stride : :obj:`int` or :obj:`tuple(int)` + padding : :obj:`int` or :obj:`tuple(int)` + dilation : :obj:`int` or :obj:`tuple(int)` + mode : :class:`CNN2D.Mode` + Direction of CNN. Either :attr:`CNN2D.Mode.Forward` or :attr:`CNN2D.Mode.Backward` + + Returns + ------- + :obj:`list(int)` + List of two integers indicating height and width of kernel. + + """ + assert(len(input_shape) == 2) + assert(len(output_shape) == 2) + if (type(stride) is int): + stride = [stride, stride] + if (type(padding) is int): + padding = [padding, padding] + if (type(dilation) is int): + dilation = [dilation, dilation] + + if (mode == CNN2D.Mode.Forward): + kern_H = (input_shape[0] + 2 * padding[0] - 1 - stride[0] * (output_shape[0] - 1)) / dilation[0] + 1 + kern_W = (input_shape[1] + 2 * padding[1] - 1 - stride[1] * (output_shape[1] - 1)) / dilation[1] + 1 + elif (mode == CNN2D.Mode.Backward): + kern_H = (output_shape[0] - (input_shape[0] - 1) * stride[0] + 2 * padding[0] - 1) / dilation[0] + 1 + kern_W = (output_shape[1] - (input_shape[1] - 1) * stride[1] + 2 * padding[1] - 1) / dilation[1] + 1 + else: + raise RuntimeError('CNN2D: Unknown mode %s!' % mode) + + if ((kern_H <= 0) or (kern_W <= 0)): + print("input shape: ", input_shape) + print("output shape: ", output_shape) + print("stride: ", stride) + print("padding: ", padding) + print("dilation: ", dilation) + print("resulting kernel size: ", [int(np.floor(kern_H)), int(np.floor(kern_W))]) + raise RuntimeError("CNN2D.compute_kernel_size: no feasible kernel size. " + "Adjust the architecture!") + + return [int(np.floor(kern_H)), int(np.floor(kern_W))] + + @classmethod + def compute_input_layer_size(cls, output_shape, kernel_size, stride, padding, dilation, mode): + """Compute input layer size that produces desired output shape with given kernel size. + + The formula is based on torch API documentation + for `Conv2d `_ + and `ConvTranspose2d `_. + + Parameters + ---------- + output_shape : :obj:`int` or :obj:`tuple(int)` + kernel_size : :obj:`int` or :obj:`tuple(int)` + stride : :obj:`int` or :obj:`tuple(int)` + padding : :obj:`int` or :obj:`tuple(int)` + dilation : :obj:`int` or :obj:`tuple(int)` + mode : :class:`CNN2D.Mode` + Direction of CNN. Either :attr:`CNN2D.Mode.Forward` or :attr:`CNN2D.Mode.Backward` + + Returns + ------- + :obj:`list(int)` + List of two integers indicating height and width of input layer. + + """ + assert(len(output_shape) == 2) + if (type(kernel_size) is int): + kernel_size = [kernel_size, kernel_size] + if (type(stride) is int): + stride = [stride, stride] + if (type(padding) is int): + padding = [padding, padding] + if (type(dilation) is int): + dilation = [dilation, dilation] + + if (mode == cls.Mode.Forward): + Hin = (output_shape[0] - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) + 1 + Win = (output_shape[1] - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) + 1 + elif (mode == cls.Mode.Backward): + Hin = (output_shape[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1) / stride[0] + 1 + Win = (output_shape[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1) / stride[1] + 1 + else: + raise RuntimeError('CNN2D: Unknown mode %s!' % mode) + + return [int(np.floor(Hin)), int(np.floor(Win))] + + @classmethod + def compute_output_layer_size(cls, input_shape, kernel_size, stride, padding, dilation, mode): + """Compute output layer size produced from given input shape and kernel size. + + The formula is based on torch API documentation + for `Conv2d `_ + and `ConvTranspose2d `_. + + Parameters + ---------- + input_shape : :obj:`int` or :obj:`tuple(int)` + kernel_size : :obj:`int` or :obj:`tuple(int)` + stride : :obj:`int` or :obj:`tuple(int)` + padding : :obj:`int` or :obj:`tuple(int)` + dilation : :obj:`int` or :obj:`tuple(int)` + mode : :class:`CNN2D.Mode` + Direction of CNN. Either :attr:`CNN2D.Mode.Forward` or :attr:`CNN2D.Mode.Backward` + + Returns + ------- + :obj:`list(int)` + List of two integers indicating height and width of output layer. + + """ + assert(len(input_shape) == 2) + if (type(kernel_size) is int): + kernel_size = [kernel_size, kernel_size] + if (type(stride) is int): + stride = [stride, stride] + if (type(padding) is int): + padding = [padding, padding] + if (type(dilation) is int): + dilation = [dilation, dilation] + + if (mode == cls.Mode.Forward): + Hout = (input_shape[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1) / stride[0] + 1 + Wout = (input_shape[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1) / stride[1] + 1 + elif (mode == cls.Mode.Backward): + Hout = (input_shape[0] - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) + 1 + Wout = (input_shape[1] - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) + 1 + else: + raise RuntimeError('CNN2D: Unknown mode %s!' % mode) + + if ((mode == cls.Mode.Forward) and ((Hout > np.floor(Hout)) or (Wout > np.floor(Wout)))): + print("input shape: ", input_shape) + print("kernel size: ", kernel_size) + print("stride: ", stride) + print("padding: ", padding) + print("dilation: ", dilation) + print("resulting output shape: ", [Hout, Wout]) + raise RuntimeError("CNN2D.compute_output_layer_size: given architecture will not return the same size backward. " + "Adjust the architecture!") + + return [int(np.floor(Hout)), int(np.floor(Wout))] + + def init_weight(self): + """Initialize weights of linear features according to Xavier uniform distribution.""" + # TODO(kevin): support other initializations? + for fc in self.fcs: + torch.nn.init.xavier_uniform_(fc.weight) + return \ No newline at end of file diff --git a/tests/test_CNN2D.py b/tests/test_CNN2D.py new file mode 100644 index 0000000..c2828a9 --- /dev/null +++ b/tests/test_CNN2D.py @@ -0,0 +1,87 @@ +from lasdi.networks import CNN2D +from random import randint +import numpy as np +import torch + +def test_set_data_shape(): + depth = 3 + layer_sizes = [[randint(2, 4), randint(45, 50), randint(45, 50)], + [randint(2, 4), randint(12, 15), randint(12, 15)], + [randint(2, 4), randint(2, 4), randint(2, 4)]] + dilations = [1] * (depth-1) + paddings = [randint(1, 3) for _ in range(depth-1)] + strides = [randint(1, 3) for _ in range(depth-1)] + + input_shape = [randint(3, 5), randint(10, 20)] + layer_sizes[0] + input_data = torch.rand(input_shape) + + cnnf = CNN2D(layer_sizes, 'forward', + strides, paddings, dilations) + + cnnf.set_data_shape(input_data.shape) + cnnf.print_data_shape() + + assert(cnnf.batch_reshape[-3:] == input_shape[-3:]) + assert(np.prod(cnnf.batch_reshape) == np.prod(input_shape)) + + output_data = cnnf(input_data) + + cnnb = CNN2D(layer_sizes[::-1], 'backward', + strides[::-1], paddings[::-1], dilations[::-1]) + cnnb.set_data_shape(input_data.shape) + cnnb.print_data_shape() + + assert(cnnb.batch_reshape == input_shape) + + input_data0 = cnnb(output_data) + assert(input_data0.shape == input_data.shape) + + return + +def test_set_data_shape2(): + depth = 3 + layer_sizes = [[1, randint(45, 50), randint(45, 50)], + [randint(2, 4), randint(12, 15), randint(12, 15)], + [randint(2, 4), randint(2, 4), randint(2, 4)]] + dilations = [1] * (depth-1) + paddings = [randint(1, 3) for _ in range(depth-1)] + strides = [randint(1, 3) for _ in range(depth-1)] + + def test_func(input_shape_): + input_data = torch.rand(input_shape_) + + cnnf = CNN2D(layer_sizes, 'forward', + strides, paddings, dilations) + + cnnf.set_data_shape(input_data.shape) + cnnf.print_data_shape() + + if ((len(input_shape_) == 2) or (input_shape_[-3] != 1)): + assert(cnnf.batch_reshape[-2:] == input_shape_[-2:]) + else: + assert(cnnf.batch_reshape[-3:] == input_shape_[-3:]) + + assert(np.prod(cnnf.batch_reshape) == np.prod(input_shape_)) + + output_data = cnnf(input_data) + + cnnb = CNN2D(layer_sizes[::-1], 'backward', + strides[::-1], paddings[::-1], dilations[::-1]) + cnnb.set_data_shape(input_data.shape) + cnnb.print_data_shape() + + assert(cnnb.batch_reshape == input_shape_) + + input_data0 = cnnb(output_data) + assert(input_data0.shape == input_data.shape) + + input_shape = layer_sizes[0][1:] + test_func(input_shape) + + input_shape = [randint(3, 5), randint(10, 20)] + layer_sizes[0][1:] + test_func(input_shape) + + input_shape = [randint(3, 5), randint(10, 20)] + layer_sizes[0] + test_func(input_shape) + + return \ No newline at end of file