Spaces:
Running on Zero
Running on Zero
File size: 1,632 Bytes
e794567 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | # -*- coding: utf-8 -*-
from dataclasses import dataclass
from transformers import PretrainedConfig
@dataclass
class VAEConfig(PretrainedConfig):
model_type: str = 'vae'
# Dataset parameters
data_path: str = ""
split_key: str = 'train'
data_aug: bool = False
img_resize: int = 1024
img_crop_size: int = 256
# Loss parameters
recon_loss_coeff: float = 1.0
kl_loss_coeff: float = 1.0
# Model parameters
in_channels: int = 1
out_channels: int = 1
num_down_blocks: int = 3
latent_channels: int = 4
vae_block_out_channels: str = '128,256,512'
# Training parameters
vae_loadcheck_path: str = ""
ft: bool = False
infer: bool = False
ifresume: bool = False
output_dir: str = ""
learning_rate: float = 1e-4
weight_decay: float = 0.0
gradient_accumulation_steps: int = 1
per_device_train_batch_size: int = 2
per_device_eval_batch_size: int = 2
num_train_epochs: int = 10
fp16: bool = False
bf16: bool = False
logging_dir: str = ""
logging_steps: int = 10
max_steps: int = -1
warmup_steps: int = 1000
save_steps: int = 1000
dataloader_num_workers: int = 8
seed: int = 6
def __init__(self, **kwargs):
# Use `super().__init__` to handle arguments from PretrainedConfig
super().__init__(**kwargs)
self.vae_block_out_channels = kwargs.get("vae_block_out_channels", self.vae_block_out_channels)
if not isinstance(self.vae_block_out_channels, list):
self.vae_block_out_channels = [int(a) for a in self.vae_block_out_channels.split(',')]
|