├── .gitignore ├── LICENSE.txt ├── README.md ├── altclip_model ├── config.json └── new_position_embedding_weight_448.pt ├── altclip_processor ├── preprocessor_config.json ├── sentencepiece.bpe.model ├── special_tokens_map.json ├── tokenizer.json └── tokenizer_config.json ├── data_process.py ├── dataset.py ├── evaluate.py ├── images ├── demo_heatmap.jpg ├── demo_image.jpg ├── model.png └── readme.md ├── model ├── basic_components.py └── model_final.py ├── requirements.txt ├── run_train.sh ├── train.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore all .idea directories and files 2 | .idea/ 3 | 4 | # Other common entries 5 | *.log 6 | *.pyc 7 | __pycache__/ 8 | .env 9 | node_modules/ 10 | dist/ 11 | .DS_Store -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # EvalMuse-Structure 2 | This repo is prepared for EvalMuse Part2-Structure Distortion Detection. 3 | Baseline for ntire 2025 'Text to Image Generation Model Quality Assessment -Track2- Structure Distortion Detection' has been released. 4 | 5 | # Framework 6 | This baseline (EM-RAHF) is inspired by paper [Rich Human Feedback for Text-to-Image Generation](https://arxiv.org/pdf/2312.10240), since the authors of RAHF did not provide code, we modified some details of the original model and achieved better performance.The details of our methods will be published in a technical report paper. 7 | ![baseline framework](images/model.png) 8 | 9 | # Train&Eval 10 | We processed the bounding boxes in the annotation file into the format of heatmap to train baseline model following RAHF 11 | | ![ori image](images/demo_image.jpg) | ![heatmap](images/demo_heatmap.jpg) | 12 | |-------------------------|-------------------------| 13 | 14 | You can process the training label into a baseline training file by executing the following script 15 | ```bash 16 | python data_process.py 17 | ``` 18 | 19 | Then you can train the baseline model by changing your own data path and training parameter settings, then run 20 | ```bash 21 | ./run_train.sh 22 | ``` 23 | note that we use the pretrained model [AltCLIP](https://huggingface.co/BAAI/AltCLIP) as vision and text encoder to achieve better performance. 24 | 25 | Evaluation and inference can be done by changing your own data path and model weight path, then run 26 | ```bash 27 | python evaluate.py 28 | ``` 29 | 30 | # Baseline Results 31 | We provide a Google Drive link for the baseline prediction result: [baseline result](https://drive.google.com/file/d/1dCnZqlSfWZbg-EVjKHefc8xuMdIlbbb1/view?usp=drive_link) 32 | The metrics score of the baseline is: 33 | 34 | | Precision | Recall | F1-score | PLCC | SROCC |Final-score | 35 | |--------------|--------------|--------------|--------------|--------------|--------------| 36 | | 0.5086 | 0.6728 | 0.5793 | 0.6945 | 0.6677 |0.6098 | 37 | 38 | # Citation and Acknowledgement 39 | 40 | If you find EvalMuse or EM-RAHF useful for your research, please consider cite our paper: 41 | ```bibtex 42 | @misc{han2024evalmuse40kreliablefinegrainedbenchmark, 43 | title={EvalMuse-40K: A Reliable and Fine-Grained Benchmark with Comprehensive Human Annotations for Text-to-Image Generation Model Evaluation}, 44 | author={Shuhao Han and Haotian Fan and Jiachen Fu and Liang Li and Tao Li and Junhui Cui and Yunqiu Wang and Yang Tai and Jingwei Sun and Chunle Guo and Chongyi Li}, 45 | year={2024}, 46 | eprint={2412.18150}, 47 | archivePrefix={arXiv}, 48 | primaryClass={cs.CV}, 49 | url={https://arxiv.org/abs/2412.18150}, 50 | } 51 | ``` 52 | For using baseline RAHF or dataset RichHF-18k, please consider cite the paper 53 | ```bibtex 54 | @inproceedings{richhf, 55 | title={Rich Human Feedback for Text-to-Image Generation}, 56 | author={Youwei Liang and Junfeng He and Gang Li and Peizhao Li and Arseniy Klimovskiy and Nicholas Carolan and Jiao Sun and Jordi Pont-Tuset and Sarah Young and Feng Yang and Junjie Ke and Krishnamurthy Dj Dvijotham and Katie Collins and Yiwen Luo and Yang Li and Kai J Kohlhoff and Deepak Ramachandran and Vidhya Navalpakkam}, 57 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, 58 | year={2024}, 59 | } 60 | ``` 61 | -------------------------------------------------------------------------------- /altclip_model/config.json: -------------------------------------------------------------------------------- 1 | { 2 | "_commit_hash": "4d06f38b304fc2a331d9f3eab77a542afafc4ffb", 3 | "_name_or_path": "BAAI/AltCLIP", 4 | "architectures": [ 5 | "AltCLIPModel" 6 | ], 7 | "direct_kd": false, 8 | "initializer_factor": 1.0, 9 | "logit_scale_init_value": 2.6592, 10 | "model_type": "altclip", 11 | "num_layers": 3, 12 | "projection_dim": 768, 13 | "text_config": { 14 | "_name_or_path": "", 15 | "add_cross_attention": false, 16 | "architectures": null, 17 | "attention_probs_dropout_prob": 0.1, 18 | "bad_words_ids": null, 19 | "begin_suppress_tokens": null, 20 | "bos_token_id": 0, 21 | "chunk_size_feed_forward": 0, 22 | "classifier_dropout": null, 23 | "cross_attention_hidden_size": null, 24 | "decoder_start_token_id": null, 25 | "diversity_penalty": 0.0, 26 | "do_sample": false, 27 | "early_stopping": false, 28 | "encoder_no_repeat_ngram_size": 0, 29 | "eos_token_id": 2, 30 | "exponential_decay_length_penalty": null, 31 | "finetuning_task": null, 32 | "forced_bos_token_id": null, 33 | "forced_eos_token_id": null, 34 | "hidden_act": "gelu", 35 | "hidden_dropout_prob": 0.1, 36 | "hidden_size": 1024, 37 | "id2label": { 38 | "0": "LABEL_0", 39 | "1": "LABEL_1" 40 | }, 41 | "initializer_factor": 0.02, 42 | "initializer_range": 0.02, 43 | "intermediate_size": 4096, 44 | "is_decoder": false, 45 | "is_encoder_decoder": false, 46 | "label2id": { 47 | "LABEL_0": 0, 48 | "LABEL_1": 1 49 | }, 50 | "layer_norm_eps": 1e-05, 51 | "length_penalty": 1.0, 52 | "max_length": 20, 53 | "max_position_embeddings": 514, 54 | "min_length": 0, 55 | "model_type": "altclip_text_model", 56 | "no_repeat_ngram_size": 0, 57 | "num_attention_heads": 16, 58 | "num_beam_groups": 1, 59 | "num_beams": 1, 60 | "num_hidden_layers": 24, 61 | "num_return_sequences": 1, 62 | "output_attentions": false, 63 | "output_hidden_states": false, 64 | "output_scores": false, 65 | "pad_token_id": 1, 66 | "pooler_fn": "cls", 67 | "position_embedding_type": "absolute", 68 | "prefix": null, 69 | "problem_type": null, 70 | "project_dim": 768, 71 | "pruned_heads": {}, 72 | "remove_invalid_values": false, 73 | "repetition_penalty": 1.0, 74 | "return_dict": true, 75 | "return_dict_in_generate": false, 76 | "sep_token_id": null, 77 | "suppress_tokens": null, 78 | "task_specific_params": null, 79 | "temperature": 1.0, 80 | "tf_legacy_loss": false, 81 | "tie_encoder_decoder": false, 82 | "tie_word_embeddings": true, 83 | "tokenizer_class": null, 84 | "top_k": 50, 85 | "top_p": 1.0, 86 | "torch_dtype": null, 87 | "torchscript": false, 88 | "transformers_version": "4.26.0.dev0", 89 | "type_vocab_size": 1, 90 | "typical_p": 1.0, 91 | "use_bfloat16": false, 92 | "use_cache": true, 93 | "vocab_size": 250002 94 | }, 95 | "text_config_dict": { 96 | "hidden_size": 1024, 97 | "intermediate_size": 4096, 98 | "num_attention_heads": 16, 99 | "num_hidden_layers": 24 100 | }, 101 | "text_model_name": null, 102 | "torch_dtype": "float32", 103 | "transformers_version": null, 104 | "vision_config": { 105 | "_name_or_path": "", 106 | "add_cross_attention": false, 107 | "architectures": null, 108 | "attention_dropout": 0.0, 109 | "bad_words_ids": null, 110 | "begin_suppress_tokens": null, 111 | "bos_token_id": null, 112 | "chunk_size_feed_forward": 0, 113 | "cross_attention_hidden_size": null, 114 | "decoder_start_token_id": null, 115 | "diversity_penalty": 0.0, 116 | "do_sample": false, 117 | "dropout": 0.0, 118 | "early_stopping": false, 119 | "encoder_no_repeat_ngram_size": 0, 120 | "eos_token_id": null, 121 | "exponential_decay_length_penalty": null, 122 | "finetuning_task": null, 123 | "forced_bos_token_id": null, 124 | "forced_eos_token_id": null, 125 | "hidden_act": "quick_gelu", 126 | "hidden_size": 1024, 127 | "id2label": { 128 | "0": "LABEL_0", 129 | "1": "LABEL_1" 130 | }, 131 | "image_size":448, 132 | "initializer_factor": 1.0, 133 | "initializer_range": 0.02, 134 | "intermediate_size": 4096, 135 | "is_decoder": false, 136 | "is_encoder_decoder": false, 137 | "label2id": { 138 | "LABEL_0": 0, 139 | "LABEL_1": 1 140 | }, 141 | "layer_norm_eps": 1e-05, 142 | "length_penalty": 1.0, 143 | "max_length": 20, 144 | "min_length": 0, 145 | "model_type": "altclip_vision_model", 146 | "no_repeat_ngram_size": 0, 147 | "num_attention_heads": 16, 148 | "num_beam_groups": 1, 149 | "num_beams": 1, 150 | "num_channels": 3, 151 | "num_hidden_layers": 24, 152 | "num_return_sequences": 1, 153 | "output_attentions": false, 154 | "output_hidden_states": false, 155 | "output_scores": false, 156 | "pad_token_id": null, 157 | "patch_size": 14, 158 | "prefix": null, 159 | "problem_type": null, 160 | "projection_dim": 512, 161 | "pruned_heads": {}, 162 | "remove_invalid_values": false, 163 | "repetition_penalty": 1.0, 164 | "return_dict": true, 165 | "return_dict_in_generate": false, 166 | "sep_token_id": null, 167 | "suppress_tokens": null, 168 | "task_specific_params": null, 169 | "temperature": 1.0, 170 | "tf_legacy_loss": false, 171 | "tie_encoder_decoder": false, 172 | "tie_word_embeddings": true, 173 | "tokenizer_class": null, 174 | "top_k": 50, 175 | "top_p": 1.0, 176 | "torch_dtype": null, 177 | "torchscript": false, 178 | "transformers_version": "4.26.0.dev0", 179 | "typical_p": 1.0, 180 | "use_bfloat16": false 181 | }, 182 | "vision_config_dict": { 183 | "hidden_size": 1024, 184 | "intermediate_size": 4096, 185 | "num_attention_heads": 16, 186 | "num_hidden_layers": 24, 187 | "patch_size": 14 188 | }, 189 | "vision_model_name": null 190 | } 191 | -------------------------------------------------------------------------------- /altclip_model/new_position_embedding_weight_448.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DYEvaLab/EvalMuse-Structure/61a36564b6ab37b7bfb8a6a6cf1d7ac870891fb2/altclip_model/new_position_embedding_weight_448.pt -------------------------------------------------------------------------------- /altclip_processor/preprocessor_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "crop_size": { 3 | "height": 224, 4 | "width": 224 5 | }, 6 | "do_center_crop": false, 7 | "do_convert_rgb": true, 8 | "do_normalize": true, 9 | "do_rescale": true, 10 | "do_resize": false, 11 | "feature_extractor_type": "CLIPFeatureExtractor", 12 | "image_mean": [ 13 | 0.48145466, 14 | 0.4578275, 15 | 0.40821073 16 | ], 17 | "image_processor_type": "CLIPImageProcessor", 18 | "image_std": [ 19 | 0.26862954, 20 | 0.26130258, 21 | 0.27577711 22 | ], 23 | "processor_class": "AltCLIPProcessor", 24 | "resample": 3, 25 | "rescale_factor": 0.00392156862745098, 26 | "size": { 27 | "shortest_edge": 224 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /altclip_processor/sentencepiece.bpe.model: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DYEvaLab/EvalMuse-Structure/61a36564b6ab37b7bfb8a6a6cf1d7ac870891fb2/altclip_processor/sentencepiece.bpe.model -------------------------------------------------------------------------------- /altclip_processor/special_tokens_map.json: -------------------------------------------------------------------------------- 1 | { 2 | "bos_token": "", 3 | "cls_token": "", 4 | "eos_token": "", 5 | "mask_token": { 6 | "content": "", 7 | "lstrip": true, 8 | "normalized": false, 9 | "rstrip": false, 10 | "single_word": false 11 | }, 12 | "pad_token": "", 13 | "sep_token": "", 14 | "unk_token": "" 15 | } 16 | -------------------------------------------------------------------------------- /altclip_processor/tokenizer_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "bos_token": "", 3 | "cls_token": "", 4 | "eos_token": "", 5 | "mask_token": { 6 | "__type": "AddedToken", 7 | "content": "", 8 | "lstrip": true, 9 | "normalized": true, 10 | "rstrip": false, 11 | "single_word": false 12 | }, 13 | "model_max_length": 512, 14 | "name_or_path": "BAAI/AltCLIP", 15 | "pad_token": "", 16 | "processor_class": "AltCLIPProcessor", 17 | "sep_token": "", 18 | "sp_model_kwargs": {}, 19 | "special_tokens_map_file": null, 20 | "tokenizer_class": "XLMRobertaTokenizer", 21 | "unk_token": "" 22 | } 23 | -------------------------------------------------------------------------------- /data_process.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import json 4 | import cv2 5 | import numpy as np 6 | ''' 7 | 将标注数据处理为baseline训练格式 train_info.pkl 8 | 注：比赛不提供验证集，若有需要请自行划分, 保存为 val_info.pkl 到同一路径 9 | ''' 10 | 11 | def overlay_heatmap_opencv(mask, image): 12 | heatmap = cv2.applyColorMap(mask.astype(np.uint8), cv2.COLORMAP_JET) 13 | overlayed_image = cv2.addWeighted(image, 0.3, heatmap, 0.7, 0) 14 | return overlayed_image 15 | 16 | train_path = 'xxx' # 训练集图片路径 17 | vis_path = 'xxx' # 热力图可视化保存路径 18 | save_path = 'xxx' # 训练数据保存路径 19 | info_path = 'train_info.json'# 训练集标注数据 20 | with open(info_path, 'r') as f: 21 | info = json.load(f) 22 | os.makedirs(save_path, exist_ok=True) 23 | 24 | vis = False 25 | if vis: 26 | os.makedirs(vis_path, exist_ok=True) 27 | data_info = {} 28 | for k,v in info.items(): 29 | try: 30 | cur_info = {} 31 | img_path = os.path.join(train_path, k+'.jpg') 32 | image = cv2.imread(img_path) 33 | height, width, _ = image.shape 34 | 35 | prompt = v['prompt_en'] 36 | score = v['mos'] 37 | bbox_info = v['bbox_info'] 38 | part_mask = np.zeros([height, width]) 39 | bbox_infos = [] 40 | 41 | for person in bbox_info: 42 | cur_mask = np.zeros([height, width]) 43 | if not person: 44 | continue 45 | else: 46 | for bbox in person: 47 | if bbox['bbox_type'] == 1: 48 | top_left, bottom_right = bbox['bbox'] 49 | cur_mask[top_left['y']:bottom_right['y'], top_left['x']:bottom_right['x']] = 1 50 | elif bbox['bbox_type'] == 2: 51 | points = bbox['bbox'] 52 | cv2.fillPoly(cur_mask, [np.array(points)], 1) 53 | else: 54 | import pdb;pdb.set_trace() 55 | part_mask = part_mask + cur_mask 56 | 57 | error_mask = part_mask.astype(np.uint8) 58 | error_mask[error_mask==1] = 64 59 | error_mask[error_mask==2] = 127 60 | error_mask[error_mask==3] = 180 61 | error_mask[error_mask==4] = 255 62 | 63 | # resize for baseline model 64 | resized_heatmap = cv2.resize(error_mask, (512, 512)) 65 | cur_info['heat_map'] = resized_heatmap 66 | cur_info['prompt'] = prompt 67 | cur_info['score'] = score 68 | data_info[k+'.jpg'] = cur_info 69 | if vis: 70 | heatmap = overlay_heatmap_opencv(error_mask, image) 71 | mask_path = os.path.join(vis_path, k+'.jpg') 72 | cv2.imwrite(mask_path, heatmap) 73 | 74 | except Exception as e: 75 | print(k,e) 76 | 77 | print(len(data_info)) 78 | with open(os.path.join(save_path, 'train_info.pkl'), 'wb') as f: 79 | pickle.dump(data_info, f) 80 | 81 | 82 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | 2 | import random 3 | import time 4 | import io 5 | import torch 6 | import pickle 7 | from transformers import AutoProcessor 8 | from torch.utils.data import DataLoader, Dataset 9 | from torchvision import transforms 10 | from PIL import Image, ImageOps 11 | from skimage.transform import resize 12 | import numpy as np 13 | 14 | random.seed(time.time()) 15 | 16 | def add_jpeg_noise(img): 17 | # Randomly add JPEG noise with quality between 70 and 100 18 | quality = random.randint(70, 100) 19 | buffer = io.BytesIO() 20 | img.save(buffer, format='JPEG', quality=quality) 21 | noisy_img = Image.open(io.BytesIO(buffer.getvalue())) 22 | return noisy_img 23 | 24 | 25 | class RAHFDataset(Dataset): 26 | def __init__(self, datapath, data_type, pretrained_processor_path, finetune=False, img_len=448): 27 | self.img_len = img_len 28 | # self.tag_word = ['human artifact', 'human mask'] # 使用siglip时需要 29 | self.tag_word = ['human artifact', 'human segmentation'] # 使用AltCLIP时需要 30 | self.finetune = finetune 31 | self.processor = AutoProcessor.from_pretrained(pretrained_processor_path) 32 | self.processor.image_processor.do_resize = False 33 | self.processor.image_processor.do_center_crop = False # 保持图片原大小 34 | self.to_tensor = transforms.ToTensor() 35 | self.datapath = datapath 36 | self.data_type = data_type 37 | # 加载pkl文件 38 | self.data_info = self.load_info() 39 | self.images = [] 40 | self.prompts_en = [] 41 | self.prompts_cn = [] 42 | self.heatmaps = [] 43 | self.scores = [] 44 | self.img_name = list(self.data_info.keys()) 45 | for i in range(len(self.img_name)): 46 | cur_img = self.img_name[i] 47 | img = Image.open(f"{self.datapath}/{self.data_type}/images/{cur_img}") 48 | self.images.append(img.resize((self.img_len, self.img_len), Image.LANCZOS)) 49 | # prompt_cn = self.data_info[cur_img]['prompt_cn'] 50 | prompt_en = self.data_info[cur_img]['prompt'] 51 | self.prompts_en.append(prompt_en) 52 | # self.prompts_cn.append(prompt_cn) 53 | 54 | artifact_map = self.data_info[cur_img]['heat_map'].astype(float) 55 | artifact_map = artifact_map/255.0 # 热力图归一化到0-1 56 | 57 | # misalignment_map = self.data_info[cur_img]['human_mask'].astype(float) # 使用0-1二值的人体mask时需要 58 | misalignment_map = np.zeros((512,512)) 59 | self.heatmaps.append([artifact_map, misalignment_map]) 60 | 61 | norm_score = (self.data_info[cur_img]['score'] - 1.0)/4.0 62 | self.scores.append((norm_score, 0)) # 人体mask没有分数 63 | if i % 1000 == 0: 64 | print(f"Processed {i} images.") 65 | 66 | if data_type == 'train' and self.finetune: 67 | self.finetune_info = self.load_info(specific_name='finetune') 68 | self.finetune_images = [] 69 | self.finetune_prompts = [] 70 | self.finetune_heatmaps = [] 71 | self.finetune_scores = [] 72 | self.finetune_img_names = list(self.finetune_info.keys()) 73 | for i in range(len(self.finetune_img_names)): 74 | cur_img = self.finetune_img_names[i] 75 | img = Image.open(f"{self.datapath}/{self.data_type}/images/{cur_img}") 76 | self.finetune_images.append(img.resize((self.img_len, self.img_len), Image.LANCZOS)) 77 | self.finetune_prompts.append(self.finetune_info[cur_img]['prompt']) 78 | artifact_map = self.finetune_info[cur_img]['artifact_map'].astype(float) 79 | misalignment_map = self.finetune_info[cur_img]['misalignment_map'].astype(float) 80 | self.finetune_heatmaps.append([artifact_map, misalignment_map]) 81 | self.finetune_scores.append((self.finetune_info[cur_img]['artifact_score'], self.finetune_info[cur_img]['misalignment_score'])) 82 | if i % 1000 == 0: 83 | print(f"Processed {i} finetuning images.") 84 | 85 | def __len__(self): 86 | return len(self.img_name) 87 | 88 | def __getitem__(self, idx): 89 | if self.data_type == 'train' and self.finetune and random.random() < 0.5: # choose finetune image to train with probability of 0.1 90 | finetune_idx = idx % len(self.finetune_img_names) 91 | finetune_img_name = self.finetune_img_names[finetune_idx] 92 | input_img = self.finetune_images[finetune_idx] 93 | input_prompt = self.finetune_prompts[finetune_idx] 94 | target_heatmaps = self.finetune_heatmaps[finetune_idx] 95 | input_img, target_heatmaps, img_pos = self.finetune_augment(input_img, target_heatmaps) 96 | cur_input = self.processor(images=input_img, text=[f"{self.tag_word[0]} {input_prompt}", f"{self.tag_word[1]} {input_prompt}"], 97 | padding="max_length", return_tensors="pt", truncation=True) 98 | 99 | cur_target = {} 100 | cur_target['artifact_map'] = (self.to_tensor(target_heatmaps[0])) 101 | cur_target['misalignment_map'] = (self.to_tensor(target_heatmaps[1])) 102 | cur_target['artifact_score'] = self.finetune_scores[finetune_idx][0] 103 | cur_target['misalignment_score'] = self.finetune_scores[finetune_idx][1] 104 | cur_target['img_name'] = finetune_img_name 105 | cur_target['img_pos'] = torch.tensor(img_pos) 106 | return cur_input, cur_target 107 | 108 | else: 109 | img_name = self.img_name[idx] 110 | input_img = self.images[idx] 111 | input_prompt = self.prompts_en[idx] 112 | target_heatmaps = self.heatmaps[idx] 113 | if self.data_type == 'train': 114 | input_img, target_heatmaps = self.data_augment(input_img, target_heatmaps) 115 | cur_input = self.processor(images=input_img, text=[f"{self.tag_word[0]} {input_prompt}", f"{self.tag_word[1]} {input_prompt}"], 116 | padding="max_length", return_tensors="pt", truncation=True) 117 | cur_target = {} 118 | cur_target['artifact_map'] = (self.to_tensor(target_heatmaps[0])) 119 | cur_target['misalignment_map'] = (self.to_tensor(target_heatmaps[1])) 120 | cur_target['artifact_score'] = self.scores[idx][0] 121 | cur_target['misalignment_score'] = self.scores[idx][1] 122 | cur_target['img_name'] = img_name 123 | cur_target['img_pos'] = torch.tensor((0,0,self.img_len)) 124 | return cur_input, cur_target 125 | 126 | def load_info(self, specific_name=None): 127 | if specific_name: 128 | print(f'Loading {specific_name} data info...') 129 | data_info = pickle.load(open(f'{self.datapath}/{specific_name}_info.pkl', 'rb')) 130 | else: 131 | print(f'Loading {self.data_type} data info...') 132 | data_info = pickle.load(open(f'{self.datapath}/{self.data_type}_info.pkl', 'rb')) 133 | return data_info 134 | 135 | def data_augment(self, img, heatmaps): 136 | 137 | if random.random() < 0.5: # 50% chance to crop 138 | crop_size = int(img.height * random.uniform(0.8, 1.0)), int(img.width * random.uniform(0.8, 1.0)) 139 | crop_region = transforms.RandomCrop.get_params(img, crop_size) 140 | img = transforms.functional.crop(img, crop_region[0], crop_region[1], crop_region[2], crop_region[3]) 141 | heatmaps = [resize(heatmap, (self.img_len, self.img_len), mode='reflect', anti_aliasing=True) 142 | for heatmap in heatmaps] 143 | heatmaps = [heatmap[crop_region[0]:crop_region[0]+crop_region[2], 144 | crop_region[1]:crop_region[1]+crop_region[3]] 145 | for heatmap in heatmaps] 146 | img = img.resize((self.img_len, self.img_len), Image.LANCZOS) 147 | heatmaps = [resize(heatmap, (512, 512), mode='reflect', anti_aliasing=True) 148 | for heatmap in heatmaps] 149 | data_transforms = transforms.Compose([ 150 | transforms.RandomApply([ 151 | transforms.ColorJitter(brightness=0.05, contrast=(0.8, 1), hue=0.025, saturation=(0.8, 1)), 152 | add_jpeg_noise 153 | ], p=0.1), 154 | transforms.RandomApply([transforms.Grayscale(3)], p=0.1) 155 | ]) 156 | 157 | img = data_transforms(img) 158 | return img, heatmaps 159 | 160 | def finetune_augment(self, img, heatmaps): 161 | 162 | data_transforms = transforms.Compose([ 163 | transforms.RandomApply([ 164 | transforms.ColorJitter(brightness=0.05, contrast=(0.8, 1), hue=0.025, saturation=(0.8, 1)), 165 | add_jpeg_noise 166 | ], p=0.2), 167 | transforms.RandomApply([transforms.Grayscale(3)], p=0.2) 168 | ]) 169 | img = data_transforms(img) 170 | # rescale gt, do nothing to heatmaps 171 | if random.random() < 0.9: # very small image 172 | scale = random.uniform(0.2, 0.5) 173 | else: 174 | scale = random.uniform(0.5, 1.0) 175 | new_len = int(scale * self.img_len) 176 | small_img = img.resize((new_len, new_len), Image.LANCZOS) 177 | top_left_x, top_left_y = random.randint(0, self.img_len-new_len), random.randint(0, self.img_len-new_len) 178 | pad_left, pad_top = top_left_x, top_left_y 179 | pad_right, pad_bottom = self.img_len - new_len - pad_left, self.img_len - new_len - pad_top 180 | pad_color = (255, 255, 255) # white padding 181 | pad_img = ImageOps.expand(small_img, border=(pad_left, pad_top, pad_right, pad_bottom), fill=pad_color) 182 | return pad_img, heatmaps, (top_left_x, top_left_y, new_len) 183 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from model.model_final import RAHF 3 | from PIL import Image 4 | import numpy as np 5 | from scipy.stats import spearmanr 6 | from dataset import RAHFDataset 7 | from torch.utils.data import DataLoader, Dataset 8 | import pickle 9 | import os 10 | import json 11 | from transformers import AutoProcessor 12 | 13 | def get_plcc_srcc(output_scores, gt_scores): 14 | # for (output_scores, gt_scores) in zip(output_scores_list, gt_scores_list): 15 | output_scores = np.array(output_scores) 16 | gt_scores = np.array(gt_scores) 17 | # Calculate PLCC (Pearson Linear Correlation Coefficient) 18 | plcc = np.corrcoef(gt_scores, output_scores)[0, 1] 19 | 20 | # Calculate SRCC (Spearman Rank Correlation Coefficient) 21 | srcc, _ = spearmanr(gt_scores, output_scores) 22 | 23 | print(f'PLCC: {plcc}') 24 | print(f'SRCC: {srcc}') 25 | 26 | def ignore_edge(heatmap): 27 | heatmap[0:5, :] = 0 # 顶部边缘 28 | heatmap[-1:-5, :] = 0 # 底部边缘 29 | heatmap[:, 0:5] = 0 # 左侧边缘 30 | heatmap[:, -1:-5] = 0 # 右侧边缘 31 | return heatmap 32 | 33 | def compute_num_params(model): 34 | import collections 35 | params = collections.defaultdict(int) 36 | bytes_per_param = 4 37 | for name, module in model.named_modules(): 38 | model_name = name.split('.')[0] 39 | if list(module.parameters()): # 只处理有参数的模块 40 | total_params = sum(p.numel() for p in module.parameters()) 41 | memory_usage_mb = (total_params * bytes_per_param) / (1024 * 1024) 42 | # print(f"模块: {name}, 参数总量: {total_params}, 显存占用: {memory_usage_mb:.2f} MB") 43 | params[model_name] += memory_usage_mb 44 | 45 | for k, v in params.items(): 46 | print(k, v,"MB") 47 | 48 | def save_pickle(obj, file_path): 49 | with open(file_path, 'wb') as f: 50 | pickle.dump(obj, f) 51 | 52 | def save_heatmap_mask(input_tensor, threshold, img_name, save_path, process_edge=False): 53 | if not os.path.exists(save_path): 54 | os.makedirs(save_path) 55 | vis_path = f'{save_path}_vis' 56 | if not os.path.exists(vis_path): 57 | os.makedirs(vis_path) 58 | input_tensor = torch.where(input_tensor > threshold, 1, 0) 59 | input_numpy = input_tensor.squeeze(0).cpu().numpy().astype(np.uint8) 60 | if process_edge: 61 | input_numpy = ignore_edge(input_numpy) 62 | vis_numpt = input_numpy * 255 63 | # Convert to PIL Image 64 | pil_image = Image.fromarray(input_numpy[0]) 65 | # Save the PIL Image 66 | pil_image.save(f"{save_path}/{img_name}.png") 67 | pil_vis = Image.fromarray(vis_numpt[0]) 68 | pil_vis.save(f"{vis_path}/{img_name}.png") 69 | 70 | 71 | def process_segment_output(outputs): 72 | normed = torch.softmax(outputs,dim=1) 73 | foreground = normed[:,1,:,:] 74 | binary_mask = (foreground>0.5).float().squeeze(0) 75 | return binary_mask 76 | 77 | def compute_badcase_detect_rate(output, target): 78 | if not output: 79 | return 0 80 | assert len(output) == len(target), "output and target must have the same length" 81 | det_count = 0 82 | for out_score, tar_score in zip(output, target): 83 | out_score = out_score*4 + 1 84 | tar_score = tar_score*4 + 1 85 | if tar_score <3 and out_score < 3: 86 | det_count += 1 87 | 88 | return det_count / len(output) 89 | 90 | def evaluate(model, dataloader, device, criterion): 91 | model.eval() 92 | loss_heatmap_im, loss_score_im, loss_heatmap_mis, loss_score_mis = 0, 0, 0, 0 93 | with torch.no_grad(): 94 | sum_heatmap_im, sum_heatmap_mis = 0.0, 0.0 95 | for inputs, targets in dataloader: 96 | inputs = inputs.to(device) 97 | 98 | outputs_im = model(inputs['pixel_values'].squeeze(1), inputs['input_ids'][:, 0, :]) # implausibility 99 | output_heatmap, target_heatmap = outputs_im[0].to(device), targets['artifact_map'].float().to(device) 100 | output_score, target_score = outputs_im[1].to(device), targets['artifact_score'].float().to(device) 101 | cur_loss_heatmap_im = criterion(output_heatmap, target_heatmap).item() 102 | loss_heatmap_im += cur_loss_heatmap_im 103 | loss_score_im += criterion(output_score, target_score).item() 104 | sum_heatmap_im += (output_heatmap * 255.0).sum().item() 105 | 106 | if targets['img_name'][0].startswith('finetune'): # check finetune data loss 107 | print(f"{targets['img_name']} artifact loss: {cur_loss_heatmap_im}") 108 | scale_factor = (255 ** 2, 4) 109 | print(f'Sum of heatmap: {sum_heatmap_im}, {sum_heatmap_mis}') 110 | return [loss_heatmap_im / len(dataloader) * scale_factor[0], loss_score_im / len(dataloader) * scale_factor[1], 111 | loss_heatmap_mis / len(dataloader), loss_score_mis / len(dataloader) * scale_factor[1]] 112 | 113 | 114 | if __name__ == '__main__': 115 | 116 | '''推理并保存计算分数的pkl文件''' 117 | gpu = "cuda:0" 118 | pretrained_processor_path = 'altclip_processor' 119 | pretrained_model_path = 'altclip_model' 120 | save_root = 'xxx' # save path of the evaluate results 121 | load_checkpoint = 'xxx' # data path of the model weight 122 | 123 | val_info_path = 'val_info.json' # data path of the val anno info 124 | val_info = json.load(open(val_info_path, 'r')) 125 | name2prompt = {k:v['prompt_en'] for k,v in val_info.items()} 126 | 127 | img_root = 'xxx' # data path of the val images 128 | img_files = os.listdir(img_root) 129 | 130 | gpu = "cuda:0" 131 | print(f'Load checkpoint {load_checkpoint}') 132 | checkpoint = torch.load(f'{load_checkpoint}', map_location='cpu') 133 | model = RAHF(pretrained_model_path=pretrained_model_path,freeze=True) 134 | model.load_state_dict(checkpoint['model']) 135 | model.cuda(gpu) 136 | model.eval() 137 | processor = AutoProcessor.from_pretrained(pretrained_processor_path) 138 | tag_word = ['human artifact', 'human segmentation'] 139 | 140 | with torch.no_grad(): 141 | preds = {} 142 | for img_file in img_files: 143 | img_name = img_file.split('.')[0] 144 | prompt = name2prompt[img_name] 145 | img_path = f'{img_root}/{img_file}' 146 | img = Image.open(img_path) 147 | image = img.resize((448, 448), Image.LANCZOS) 148 | 149 | cur_input = processor(images=image, text=[f"{tag_word[0]} {prompt}", f"{tag_word[1]} {prompt}"], 150 | padding="max_length", return_tensors="pt", truncation=True) 151 | inputs_pixel_values, inputs_ids_im = cur_input['pixel_values'].to(gpu), cur_input['input_ids'][0, :].unsqueeze(0).to(gpu) 152 | 153 | heatmap, score = model(inputs_pixel_values, inputs_ids_im, need_score=True) 154 | print(f'heatmap: {heatmap.shape}, score: {score}') 155 | 156 | ori_heatmap = torch.round(heatmap * 255.0) 157 | heatmap_treshold = 40 158 | input_tensor = torch.where(ori_heatmap > heatmap_treshold, 1, 0) 159 | saved_output_im_map = input_tensor.squeeze(0).cpu().numpy().astype(np.uint8) 160 | preds[img_name[:-4]] = { 161 | "score":score.item(), 162 | "pred_area": saved_output_im_map 163 | } 164 | 165 | with open(f'{save_root}/baseline_results.pkl', 'wb') as f: 166 | pickle.dump(preds, f) 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | -------------------------------------------------------------------------------- /images/demo_heatmap.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DYEvaLab/EvalMuse-Structure/61a36564b6ab37b7bfb8a6a6cf1d7ac870891fb2/images/demo_heatmap.jpg -------------------------------------------------------------------------------- /images/demo_image.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DYEvaLab/EvalMuse-Structure/61a36564b6ab37b7bfb8a6a6cf1d7ac870891fb2/images/demo_image.jpg -------------------------------------------------------------------------------- /images/model.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DYEvaLab/EvalMuse-Structure/61a36564b6ab37b7bfb8a6a6cf1d7ac870891fb2/images/model.png -------------------------------------------------------------------------------- /images/readme.md: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /model/basic_components.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Attention(nn.Module): 6 | def __init__(self, dim): 7 | super(Attention, self).__init__() 8 | self.to_q = nn.Linear(dim, dim, bias=False) 9 | self.to_k = nn.Linear(dim, dim, bias=False) 10 | self.to_v = nn.Linear(dim, dim, bias=False) 11 | self.attend = nn.Softmax(dim=-1) 12 | self.to_out = nn.Linear(dim, dim, bias=False) 13 | 14 | def forward(self, x): 15 | q = self.to_q(x) 16 | k = self.to_k(x) 17 | v = self.to_v(x) 18 | 19 | attn = self.attend(torch.matmul(q, k.transpose(-2, -1)) / (x.shape[-1] ** 0.5)) 20 | out = torch.matmul(attn, v) 21 | return self.to_out(out) 22 | 23 | class FeedForward(nn.Module): 24 | def __init__(self, dim, hidden_dim): 25 | super(FeedForward, self).__init__() 26 | self.ff = nn.Sequential( 27 | nn.Linear(dim, hidden_dim), 28 | nn.GELU(), 29 | nn.Linear(hidden_dim, dim) 30 | ) 31 | 32 | def forward(self, x): 33 | return self.ff(x) 34 | 35 | class Residual(nn.Module): 36 | def __init__(self, module): 37 | super(Residual, self).__init__() 38 | self.module = module 39 | 40 | def forward(self, x): 41 | return x + self.module(x) -------------------------------------------------------------------------------- /model/model_final.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """RAHF_model.ipynb 3 | 4 | Automatically generated by Colab. 5 | 6 | Original file is located at 7 | https://colab.research.google.com/drive/1CQSm0-C9TBWVDIHHyqgtfpoYC42esbTs 8 | """ 9 | 10 | import torch 11 | import torch.nn as nn 12 | 13 | from transformers import AutoProcessor, AutoModel, AutoConfig 14 | 15 | from .basic_components import Residual, Attention, FeedForward 16 | import torch.nn.functional as F 17 | 18 | 19 | class VisionTransformer(nn.Module): 20 | def __init__(self, pretrained_model, freeze, new_position_embedding_weight): 21 | super(VisionTransformer, self).__init__() 22 | self.ViT = pretrained_model.vision_model 23 | self.new_position_embedding_weight = new_position_embedding_weight 24 | # 插值vit embeddings 25 | # self.interpolate_embeddings() 26 | self.load_interpolate_embeddings() 27 | 28 | self.freeze = freeze 29 | if self.freeze: 30 | for param in self.ViT.parameters(): 31 | param.requires_grad = False 32 | else: 33 | for param in self.ViT.post_layernorm.parameters(): 34 | param.requires_grad = False 35 | 36 | def forward(self, x): 37 | x = self.ViT(x) 38 | return x 39 | 40 | def unfreeze(self): 41 | for param in self.ViT.parameters(): 42 | param.requires_grad = True 43 | for param in self.ViT.post_layernorm.parameters(): 44 | param.requires_grad = False 45 | # 使用siglip时需要 46 | # for param in self.ViT.head.parameters(): # unused 47 | # param.requires_grad = False 48 | 49 | def load_interpolate_embeddings(self): 50 | with torch.no_grad(): 51 | self.ViT.embeddings.position_embedding.weight.copy_(self.new_position_embedding_weight) 52 | 53 | def interpolate_embeddings(self): 54 | self.ViT.embeddings.position_ids = torch.arange(1025).unsqueeze(0) 55 | position_embedding = self.ViT.embeddings.position_embedding.weight 56 | cls_pos,ori_pos = position_embedding[0], position_embedding[1:] 57 | ori_pos = ori_pos.view(16,16,-1).unsqueeze(0).permute(0,3,1,2) 58 | 59 | resized_pos = F.interpolate(ori_pos, size=(32, 32), mode='bilinear') 60 | resized_pos = resized_pos.squeeze(0).permute(1,2,0).view(1024,-1) 61 | new_position_embedding_weight = torch.cat((cls_pos.unsqueeze(0),resized_pos),dim=0) 62 | new_position_embedding = nn.Embedding(1025, 1024) 63 | with torch.no_grad(): 64 | new_position_embedding.weight.copy_(new_position_embedding_weight) 65 | self.ViT.embeddings.position_embedding = new_position_embedding 66 | 67 | class TextEmbedding(nn.Module): 68 | def __init__(self, pretrained_model, freeze): 69 | super(TextEmbedding, self).__init__() 70 | 71 | # AltCLIP 72 | self.text_embedding = pretrained_model.text_model.roberta.embeddings 73 | self.freeze = freeze 74 | if self.freeze: 75 | for param in self.text_embedding.parameters(): 76 | param.requires_grad = False 77 | 78 | def forward(self, x): 79 | x = self.text_embedding(x) 80 | return x 81 | 82 | def unfreeze(self): 83 | for param in self.text_embedding.parameters(): 84 | param.requires_grad = True 85 | 86 | 87 | 88 | class LayerPair(nn.Module): 89 | def __init__(self, dim, hidden_dim): 90 | super(LayerPair, self).__init__() 91 | self.norm1 = nn.LayerNorm(dim) 92 | self.attention = Residual(Attention(dim)) 93 | self.norm2 = nn.LayerNorm(dim) 94 | self.feed_forward = Residual(FeedForward(dim, hidden_dim)) 95 | 96 | def forward(self, x): 97 | x = self.norm1(x) 98 | x = self.attention(x) 99 | x = self.norm2(x) 100 | x = self.feed_forward(x) 101 | return x 102 | 103 | class SelfAttention(nn.Module): 104 | def __init__(self, num_layers=6, dim=768, hidden_dim=2048): 105 | super(SelfAttention, self).__init__() 106 | self.layers = nn.ModuleList([LayerPair(dim, hidden_dim) for _ in range(num_layers)]) 107 | self.norm = nn.LayerNorm(dim) 108 | 109 | def forward(self, x): 110 | for layer in self.layers: 111 | x = layer(x) 112 | return self.norm(x) 113 | 114 | class HeatmapPredictor(nn.Module): 115 | def __init__(self, conv_info = [768, 384, 384], deconv_info=[384, 768, 384, 384, 192]): 116 | super(HeatmapPredictor, self).__init__() 117 | self.filter_size = deconv_info 118 | self.conv_layers = nn.Sequential( 119 | nn.Conv2d(in_channels=conv_info[0], out_channels=conv_info[1], kernel_size=(3, 3), stride=(1, 1), padding=1), 120 | nn.LayerNorm([conv_info[1], 32, 32]), 121 | nn.ReLU(), 122 | nn.Conv2d(in_channels=conv_info[1], out_channels=conv_info[2], kernel_size=(3, 3), stride=(1, 1), padding=1), 123 | nn.LayerNorm([conv_info[2], 32, 32]), 124 | nn.ReLU() 125 | ) 126 | self.deconv_layers = nn.ModuleList([ 127 | nn.ModuleList([ 128 | nn.ConvTranspose2d(in_channels=self.filter_size[i], 129 | out_channels=self.filter_size[i+1], 130 | kernel_size=(3, 3), 131 | stride=(2, 2), 132 | padding=1, output_padding=1), 133 | nn.LayerNorm([self.filter_size[i+1], 32*2**(i+1), 32*2**(i+1)]), 134 | nn.ReLU(), 135 | nn.Conv2d(in_channels=self.filter_size[i+1], 136 | out_channels=self.filter_size[i+1], 137 | kernel_size=(3, 3), stride=(1, 1), padding=1), 138 | nn.LayerNorm([self.filter_size[i+1], 32*2**(i+1), 32*2**(i+1)]), 139 | nn.ReLU(), 140 | nn.Conv2d(in_channels=self.filter_size[i+1], 141 | out_channels=self.filter_size[i+1], 142 | kernel_size=(3, 3), stride=(1, 1), padding=1), 143 | nn.LayerNorm([self.filter_size[i+1], 32*2**(i+1), 32*2**(i+1)]), 144 | nn.ReLU(), 145 | ]) for i in range(len(self.filter_size)-1) 146 | ]) 147 | self.final_layers = nn.Sequential( 148 | nn.Conv2d(in_channels=self.filter_size[-1], out_channels=1, kernel_size=(3, 3), stride=(1, 1), padding=1), 149 | nn.Sigmoid() 150 | ) 151 | 152 | def forward(self, x): 153 | x = self.conv_layers(x) 154 | for layer in self.deconv_layers: 155 | for deconv in layer: 156 | x = deconv(x) 157 | x = self.final_layers(x) 158 | return x 159 | 160 | class ScorePredictor(nn.Module): 161 | def __init__(self, filter_info=[768, 768, 384, 128, 64]): 162 | super(ScorePredictor, self).__init__() 163 | self.filter_size = filter_info 164 | self.conv_layers = nn.ModuleList([ 165 | nn.ModuleList([ 166 | nn.Conv2d(in_channels=self.filter_size[i], 167 | out_channels=self.filter_size[i + 1], 168 | kernel_size=(2, 2), stride=(1, 1)), 169 | nn.LayerNorm([self.filter_size[i + 1], 32-(i+1), 32-(i+1)]), 170 | nn.ReLU(), 171 | ]) for i in range(len(self.filter_size) - 1) 172 | ]) 173 | self.flatten_size = self.filter_size[-1] * (32-(len(self.filter_size)-1))**2 174 | self.fc_layers = nn.Sequential( 175 | nn.Linear(in_features=self.flatten_size, out_features=2048), 176 | nn.ReLU(), 177 | nn.Linear(in_features=2048, out_features=1024), 178 | nn.ReLU(), 179 | nn.Linear(in_features=1024, out_features=1), 180 | nn.Sigmoid() 181 | ) 182 | 183 | def forward(self, x): 184 | for layers in self.conv_layers: 185 | for layer in layers: 186 | x = layer(x) 187 | x = torch.flatten(x, start_dim=1) 188 | x = self.fc_layers(x) 189 | return x 190 | 191 | class RAHF(nn.Module): 192 | def __init__(self, pretrained_model_path, freeze): 193 | super(RAHF, self).__init__() 194 | # interpolate = True if 'altclip' in pretrained_model_path else False 195 | pretrained_config = AutoConfig.from_pretrained(pretrained_model_path) 196 | pretrained_config.vision_config.image_size = 448 197 | pretrained_model = AutoModel.from_pretrained(pretrained_model_path, config=pretrained_config, ignore_mismatched_sizes=True) 198 | new_position_embedding_weight = torch.load(f"{pretrained_model_path}/new_position_embedding_weight_448.pt") 199 | self.image_encoder = VisionTransformer(pretrained_model, freeze, new_position_embedding_weight) 200 | self.text_encoder = TextEmbedding(pretrained_model, freeze) 201 | self.self_attention = SelfAttention(dim=1024, hidden_dim=2048) 202 | self.heatmap_predictor = HeatmapPredictor(conv_info = [1024, 512, 512], deconv_info=[512, 1024, 512, 512, 256]) 203 | self.score_predictor = ScorePredictor(filter_info=[1024, 1024, 512, 128, 64]) 204 | 205 | def forward(self, image, prompt, need_score=True): 206 | image_token = self.image_encoder(image).last_hidden_state 207 | text_token = self.text_encoder(prompt) 208 | x = torch.cat([image_token, text_token], dim=1) 209 | x = self.self_attention(x) 210 | feature_map = x[:, 1:1025, :].clone().view(-1, 32, 32, 1024).permute(0, 3, 1, 2) 211 | heatmap = self.heatmap_predictor(feature_map) 212 | 213 | if not need_score: 214 | return heatmap 215 | else: 216 | score = self.score_predictor(feature_map) 217 | return heatmap, score 218 | 219 | if __name__ == "__main__": 220 | model = RAHF() -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # torch 2 | transformers==4.40.1 3 | pillow 4 | sentencepiece 5 | protobuf 6 | scikit-image 7 | matplotlib 8 | opencv-python-headless 9 | numpy 10 | scipy 11 | skimage -------------------------------------------------------------------------------- /run_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Change directory to the target location 4 | # cd /mnt/bn/rahf/mlx/users/jincheng.liang/repo/12094/RAHF || exit 5 | 6 | # Print the environment variable to verify it is set 7 | # echo "PYTORCH_CUDA_ALLOC_CONF is set to $PYTORCH_CUDA_ALLOC_CONF" 8 | 9 | # Install Python dependencies from requirements.txt 10 | pip3 install -r requirements.txt 11 | 12 | # Force reinstall numpy to a specific version 13 | pip3 install --force-reinstall numpy==1.25.2 14 | 15 | # Set environment variable for CUDA memory allocation 16 | # export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512 17 | 18 | # Run the Python script 19 | # python3 train_public_cluster.py 20 | python3 -m torch.distributed.launch \ 21 | --nproc_per_node=8 \ 22 | --nnodes=1 \ 23 | --node_rank=0 \ 24 | --master_addr="localhost" \ 25 | --master_port=12345 \ 26 | train.py \ 27 | --experiment_name "xxx" \ 28 | --lr 2e-5 \ 29 | --iters 2000 \ 30 | --batch_size 4 \ 31 | --accumulate_step 8 \ 32 | --val_iter 50 \ 33 | --save_iter 100 \ 34 | --warmup \ 35 | --data_path xxx \ 36 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """RAHF_train.ipynb 3 | 4 | Automatically generated by Colab. 5 | 6 | Original file is located at 7 | https://colab.research.google.com/drive/1n8Bug-l4fVCAXA7kLDJmhbKkN52jIsXm 8 | """ 9 | import os 10 | import time 11 | import torch 12 | from model.model_final import RAHF 13 | from torch.utils.data import DataLoader 14 | from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR 15 | import torch.nn as nn 16 | import logging 17 | import argparse 18 | import torch.distributed as dist 19 | from dataset import RAHFDataset 20 | from utils import print_log_iter, eval_in_training, save_in_training, final_save 21 | 22 | def train(args): 23 | try: 24 | local_rank = int(os.environ['LOCAL_RANK']) 25 | except: 26 | local_rank = dist.get_rank() 27 | torch.cuda.set_device(local_rank) 28 | gpu = f'cuda:{local_rank}' 29 | print(f'GPU: {gpu}') 30 | torch.cuda.empty_cache() 31 | 32 | save_path = f'{args.bytenas_path}/experiments/{args.experiment_name}' 33 | if not os.path.exists(save_path): 34 | os.makedirs(save_path) 35 | logging.basicConfig(filename=f'{save_path}/{args.experiment_name}.log', level=logging.INFO, format='%(asctime)s - %(message)s') 36 | logger = logging.getLogger() 37 | datapath = args.data_path 38 | print('datapath', datapath) 39 | print('bytenas path', args.bytenas_path) 40 | pretrained_processor_path = 'altclip_processor' 41 | pretrained_model_path = 'altclip_model' 42 | 43 | dist.init_process_group(backend='nccl', init_method='env://') 44 | args.rank = dist.get_rank() 45 | 46 | print(f'Using {torch.cuda.device_count()} GPUs') 47 | print(f'Freeze the pretrained componenets? {args.warmup}. Preparing model...') 48 | model = RAHF(pretrained_model_path=pretrained_model_path,freeze=args.warmup) 49 | model.cuda(gpu) 50 | if len(args.load_checkpoint) > 0: 51 | load_checkpoint = f'{args.bytenas_path}/experiments/{args.load_checkpoint}' 52 | print(f'Load checkpoint {load_checkpoint}') 53 | checkpoint = torch.load(f'{load_checkpoint}', map_location='cpu') 54 | model.load_state_dict(checkpoint['model']) 55 | else: 56 | print('Train from scratch') 57 | model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 58 | # model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu], broadcast_buffers=False, find_unused_parameters=True) 59 | model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu], find_unused_parameters=False) 60 | optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) 61 | # lr_lambda = lambda step: min((step+1) / 500.0, 1.0) 62 | # lr_lambda = lambda step: min(1.0/math.sqrt(step+1), 1.0) 63 | # lr_lambda = lambda step: 1.0 64 | # scheduler = CyclicLR(optimizer, base_lr=1e-5, max_lr=1e-3, step_size_up=400, cycle_momentum=False) 65 | # scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda) 66 | scheduler = CosineAnnealingLR(optimizer, T_max=args.iters, eta_min=args.min_lr, last_epoch=-1) 67 | if len(args.load_checkpoint) > 0: 68 | optimizer.load_state_dict(checkpoint['optimizer']) 69 | scheduler.load_state_dict(checkpoint['scheduler']) 70 | 71 | criterion = torch.nn.MSELoss().to(gpu) 72 | def criterion_heatmap(output_heatmap, target_heatmap, weighted_loss, criterion): 73 | if weighted_loss: 74 | mse_heatmap = torch.nn.MSELoss(reduction='none').to(gpu) 75 | loss_heatmap = mse_heatmap(output_heatmap, target_heatmap) 76 | loss_weights_heatmap = target_heatmap + 1.0 / 255.0 # loss weight related to pixel value, prevent from output all 0 77 | weighted_loss = (loss_heatmap * loss_weights_heatmap).sum(dim=(-2, -1)) / loss_weights_heatmap.sum(dim=(-2, -1)) 78 | weighted_loss = weighted_loss.sum() / weighted_loss.shape[0] 79 | return weighted_loss 80 | else: 81 | return criterion(output_heatmap, target_heatmap) 82 | if args.weighted_loss: 83 | print('Use weighted MSE loss') 84 | else: 85 | print('Use normal MSE loss') 86 | 87 | def train_loop(model, train_dataloader, val_dataloader, iter_counter, epoch_counter, end_iter, accumulate_step): 88 | print(f"iter:{iter_counter}, epoch:{epoch_counter}, end:{end_iter}, accumlate:{accumulate_step}") 89 | while True: 90 | model.train() 91 | print(f'Epoch {epoch_counter}') 92 | train_dataloader.sampler.set_epoch(epoch_counter) 93 | iter_loss = [[], [], [], []] 94 | for batch_id, (inputs, targets) in enumerate(train_dataloader): 95 | inputs = inputs.to(gpu) 96 | inputs_pixel_values, inputs_ids_im, inputs_ids_mis = inputs['pixel_values'].squeeze(1), inputs['input_ids'][:, 0, :], inputs['input_ids'][:, 1, :] 97 | outputs_im = model(inputs_pixel_values, inputs_ids_im, need_score=True) # implausibility 98 | # implausibility heatmap 99 | output_heatmap, target_heatmap = outputs_im[0].to(gpu), targets['artifact_map'].float().to(gpu) 100 | loss_im = criterion_heatmap(output_heatmap, target_heatmap, args.weighted_loss, criterion) 101 | 102 | # implausibility score 103 | output_score, target_score = outputs_im[1].to(gpu), targets['artifact_score'].float().to(gpu) 104 | 105 | loss_score = criterion(output_score, target_score) 106 | # implausibility loss 107 | iter_loss[0].append(loss_im.item()) 108 | iter_loss[1].append(loss_score.item()) 109 | loss_im = loss_im + loss_score 110 | loss_im.backward() 111 | 112 | iter_loss[2].append(0) 113 | iter_loss[3].append(0) 114 | 115 | if (batch_id + 1) % accumulate_step == 0 or batch_id == len(train_dataloader): 116 | optimizer.step() 117 | optimizer.zero_grad() 118 | scheduler.step() 119 | iter_counter += 1 120 | dist.barrier() 121 | print_log_iter(optimizer, iter_counter, iter_loss, logger) 122 | iter_loss = [[], [], [], []] 123 | if iter_counter % args.val_iter == 0: 124 | eval_in_training(model, val_dataloader, gpu, criterion, iter_counter, logger) 125 | if iter_counter % args.save_iter == 0: 126 | save_in_training(model, optimizer, scheduler, iter_counter, save_path) 127 | if iter_counter >= end_iter: 128 | return iter_counter, epoch_counter 129 | 130 | epoch_counter += 1 131 | 132 | print('Preparing dataloader...') 133 | train_dataset = RAHFDataset(datapath, 'train', pretrained_processor_path, finetune=False, img_len=448) 134 | train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) 135 | train_dataloader = DataLoader(dataset=train_dataset, 136 | batch_size=args.batch_size, 137 | shuffle=False, 138 | num_workers=8, 139 | pin_memory=True, 140 | sampler=train_sampler) 141 | 142 | val_dataset = RAHFDataset(datapath, 'val', pretrained_processor_path, img_len=448) 143 | val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset) 144 | val_dataloader = DataLoader(dataset=val_dataset, 145 | batch_size=args.batch_size, 146 | # batch_size=1, # to get finetune performance 147 | shuffle=False, 148 | pin_memory=True, 149 | num_workers=8, 150 | sampler=val_sampler) 151 | 152 | train_dataloader1 = DataLoader(dataset=train_dataset, 153 | batch_size=1, 154 | shuffle=False, 155 | num_workers=8, 156 | pin_memory=True, 157 | sampler=train_sampler) 158 | 159 | dist.barrier() 160 | print('Training...') 161 | iter_counter = 0 162 | epoch_counter = 0 163 | start_time = time.time() 164 | torch.autograd.set_detect_anomaly(True) 165 | iter_counter, epoch_counter = train_loop(model, train_dataloader, val_dataloader, iter_counter, epoch_counter, args.iters//2, args.accumulate_step) 166 | dist.barrier() 167 | model.module.image_encoder.unfreeze() 168 | model.module.text_encoder.unfreeze() 169 | print('Unfreeze image encoder and text encoder after 1000 iterations.') 170 | del train_dataloader 171 | iter_counter, epoch_counter = train_loop(model, train_dataloader1, val_dataloader, iter_counter, epoch_counter, args.iters, 32) 172 | dist.barrier() 173 | final_save(model, optimizer, scheduler, start_time, save_path) 174 | dist.destroy_process_group() 175 | 176 | def main(): 177 | parser = argparse.ArgumentParser() 178 | # Training settings 179 | # parser.add_argument('-gpu_n', default=4, type=int, help="how many gpu") 180 | # parser.add_argument('-g', '--gpuid', default=0, type=int, help="which gpu to use") 181 | parser.add_argument("--local-rank", default=0, type=int, help='rank in current node') 182 | # Experiment settings 183 | parser.add_argument("--experiment_name", required=True, type=str, help="name of this experiment") 184 | parser.add_argument("--load_checkpoint", default='', type=str, help="the name of the checkpoint to be loaded") 185 | parser.add_argument("--bytenas_path", type=str, default='xxx', help="path of bytenas") # 存放实验相关内容 186 | parser.add_argument("--data_path", type=str, default='xxx', help="path of data") # 训练/测试数据存放路径 187 | parser.add_argument('--iters', required=True, type=int, metavar='N', help='number of total iterations to run') 188 | parser.add_argument('--batch_size', default=4, type=int, metavar='N', help='the batchsize for each gpu') 189 | parser.add_argument('--accumulate_step', default=16, type=int, metavar='N', help='accumulate_step * batch_size = actual batch size') 190 | parser.add_argument('--lr', default=1e-3, type=float, help='base learning rate') 191 | parser.add_argument('--min_lr', default=0.0, type=float, help='min learning rate') 192 | parser.add_argument('--val_iter', default=25, type=int, metavar='N', help='number of iterations to run validation') 193 | parser.add_argument('--save_iter', default=200, type=int, metavar='N', help='number of iterations to save') 194 | parser.add_argument('--warmup', action='store_true', help='whether to freeze the pretrained components') 195 | parser.add_argument('--weighted_loss', action='store_true', help='weighted loss for heatmap prevent output all 0') 196 | # parser.add_argument('--loss_weights', default=[1.0, 1.0, 0.5, 0.5], help='loss weight for: implausibility heatmap & score, misalignment heatmap & score') 197 | args = parser.parse_args() 198 | train(args) 199 | 200 | if __name__ == '__main__': 201 | main() -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | from functools import wraps 4 | import torch.distributed as dist 5 | from evaluate import evaluate 6 | 7 | def only_rank_0(func): 8 | @wraps(func) 9 | def wrapper(*args, **kwargs): 10 | if dist.get_rank() == 0: 11 | return func(*args, **kwargs) 12 | return wrapper 13 | 14 | @only_rank_0 15 | def print_log_iter(optimizer, iter_counter, iter_loss, logger): 16 | current_lr = optimizer.param_groups[0]['lr'] 17 | aver_loss_im_heatmap = sum(iter_loss[0]) / len(iter_loss[0]) 18 | aver_loss_im_score = sum(iter_loss[1]) / len(iter_loss[1]) 19 | aver_loss_sematic_heatmap = sum(iter_loss[2]) / len(iter_loss[2]) 20 | # aver_loss_mis_score = sum(iter_loss[3]) / len(iter_loss[3]) 21 | aver_loss_mis_heatmap,aver_loss_mis_score = 0,0 22 | logger.info(f"Iteration {iter_counter}: Learning Rate = {current_lr}\n" 23 | f"Implausibility Heatmap Loss = {aver_loss_im_heatmap}, Implausibility Score Loss = {aver_loss_im_score}\n" 24 | f"Sematic Heatmap Loss = {aver_loss_sematic_heatmap}, Misalignment Score Loss = {aver_loss_mis_score}") 25 | print(f"Iteration {iter_counter}: Learning Rate = {current_lr}\n" 26 | f"Implausibility Heatmap Loss = {aver_loss_im_heatmap}, Implausibility Score Loss = {aver_loss_im_score}\n" 27 | f"Sematic Heatmap Loss = {aver_loss_sematic_heatmap}, Misalignment Score Loss = {aver_loss_mis_score}") 28 | 29 | @only_rank_0 30 | def eval_in_training(model, val_dataloader, device, criterion, iter_counter, logger): 31 | # val_loss = evaluate(model=model.module, dataloader=val_dataloader, device=device, criterion=criterion,criterion2=criterion2) 32 | val_loss = evaluate(model=model.module, dataloader=val_dataloader, device=device, criterion=criterion) 33 | logger.info(f"Iteration {iter_counter} Validation:\n" 34 | f"Implausibility Heatmap Loss = {val_loss[0]}, Implausibility Score Loss = {val_loss[1]}\n" 35 | f"Sematic Heatmap Loss = {val_loss[2]}, Misalignment Score Loss = {val_loss[3]}") 36 | print(f"Iteration {iter_counter} Validation:\n" 37 | f"Implausibility Heatmap Loss = {val_loss[0]}, Implausibility Score Loss = {val_loss[1]}\n" 38 | f"Sematic Heatmap Loss = {val_loss[2]}, Misalignment Score Loss = {val_loss[3]}") 39 | 40 | @only_rank_0 41 | def save_in_training(model, optimizer, scheduler, iter_counter, save_path): 42 | checkpoint = { 43 | 'model': model.module.state_dict(), 44 | 'optimizer': optimizer.state_dict(), 45 | 'scheduler': scheduler.state_dict() 46 | } 47 | torch.save(checkpoint, f'{save_path}/{iter_counter}.pth') 48 | print(f"Model weights saved to {save_path}/{iter_counter}.pth") 49 | 50 | @only_rank_0 51 | def final_save(model, optimizer, scheduler, start_time, save_path): 52 | end_time = time.time() 53 | checkpoint = { 54 | 'model': model.module.state_dict(), 55 | 'optimizer': optimizer.state_dict(), 56 | 'scheduler': scheduler.state_dict() 57 | } 58 | total_minutes = (end_time - start_time) / 60 59 | torch.save(checkpoint, f'{save_path}/last.pth') 60 | print(f"Model weights saved to {save_path}/last.pth. Total training time: {total_minutes:.2f} minutes") --------------------------------------------------------------------------------