Code for paper Micrometer: Micromechanics Transformer for Predicting Mechanical Responses of Heterogeneous Materials.
Our code has been tested with a Linux environment using the following configuration:
- Python 3.9
- CUDA 12.4
- CUDNN 8.9
- JAX 0.4.26
First, clone the repository:
git clone https://github.com/PredictiveIntelligenceLab/micrometer.git
cd micrometer
Then, install the required packages
pip3 install -U pip
pip3 install --upgrade jax jaxlib
pip3 install --upgrade -r requirements.txt
Finally, install the package:
pip3 install -e .
Our dataset can be downloaded from the following Google Drive links:
| Name | Link |
|---|---|
| CMME | Link |
| Homogenization | Link |
| Multiscale Modelling | Link |
| Transfer Learning | Link |
First, please place the downloaded the dataset and change the data path in
the corresponding config file, e.g. in configs/base.py:
dataset.data_path = "path/to/dataset"
Then, to train our Micrometer (e.g. cvit_b_16), run the following command:
python3 main.py --config=configs/base.py:cvit_b_16
The user can also train other models by changing the model name
in the above command, e.g. cvit_b_16 to cvit_l_8. We also provide UNet or FNO of different configurations as the backbone model,
which can be found in configs/models.py.
To specify the GPU device, for example, use the following command:
CUDA_VISIBLE_DEVICES=0,1,2,3 python3 main.py --config=configs/base.py:cvit_b_16
To evaluate the trained model (e.g. cvit_b_16), run the following command:
python3 main.py --config=configs/base.py:cvit_b_16 --config.mode=eval
With the pre-trained model, we can perform homogenization by running the following command:
python3 main.py --config=configs/homo.py:cvit_b_16
With the pre-trained model, we can perform multiscale modelling
python3 main.py --config=configs/multiscale.py:cvit_b_16
We can also fine-tune our pretrained model on new datasets, which is configured in
finetune_vf.py or finetune_ch.py
python3 main.py --config=configs/finetune_ch.py:cvit_b_16