| from transformers import PretrainedConfig |
| from typing import List |
|
|
| class WNet3DConfig(PretrainedConfig): |
| model_type = "WNet" |
| def __init__( |
| self, |
| in_ch=1, |
| out_ch=5, |
| init_features=64, |
| **kwargs): |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.init_features = init_features |
| super().__init__(**kwargs) |
|
|
| class AttWNet3DConfig(PretrainedConfig): |
| model_type = "AttWNet" |
| def __init__( |
| self, |
| in_ch=1, |
| out_ch=5, |
| init_features=64, |
| **kwargs): |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.init_features = init_features |
| super().__init__(**kwargs) |
|
|
| class WNetMSS3DConfig(PretrainedConfig): |
| model_type = "WNetMSS" |
| def __init__( |
| self, |
| in_ch=1, |
| out_ch=5, |
| init_features=64, |
| **kwargs): |
| self.in_ch = in_ch |
| self.out_ch = out_ch |
| self.init_features = init_features |
| super().__init__(**kwargs) |
