| from transformers import PretrainedConfig | |
| from typing import List | |
| class UNetFTSRConfig(PretrainedConfig): | |
| model_type = "UNet" | |
| def __init__( | |
| self, | |
| in_channels=1, | |
| n_classes=1, | |
| depth=3, | |
| wf=6, | |
| padding=True, | |
| batch_norm=False, | |
| up_mode='upconv', | |
| dropout=False, | |
| **kwargs): | |
| self.in_channels = in_channels | |
| self.n_classes = n_classes | |
| self.depth = depth | |
| self.wf = wf | |
| self.padding = padding | |
| self.batch_norm = batch_norm | |
| self.up_mode = up_mode | |
| self.dropout = dropout | |
| super().__init__(**kwargs) | |