| from transformers import PretrainedConfig |
| from typing import List |
|
|
| class UNet3DConfig(PretrainedConfig): |
| model_type = "UNet" |
| def __init__( |
| self, |
| in_ch=1, |
| out_ch=1, |
| init_features=64, |
| **kwargs): |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.init_features = init_features |
| super().__init__(**kwargs) |
|
|
| class UNetMSS3DConfig(PretrainedConfig): |
| model_type = "UNetMSS" |
| def __init__( |
| self, |
| in_ch=1, |
| out_ch=1, |
| init_features=64, |
| **kwargs): |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.init_features = init_features |
| super().__init__(**kwargs) |
