[refactor] configuration default values

Former-commit-id: d7ca7685d6a8c6d015082dda241bc98ab156770d
This commit is contained in:
qqingzheng
2024-03-06 15:59:13 +00:00
parent 4549eeeab3
commit 7a78ef2a97
2 changed files with 15 additions and 24 deletions
@@ -1,17 +1,17 @@
from ..configuration_videobase import VideoBaseConfiguration
from dataclasses import dataclass
from typing import Tuple
from dataclasses import dataclass, field
from typing import Tuple, Union
@dataclass
class VQVAEConfiguration:
embedding_dim: int
n_codes: int
n_hiddens: int
n_res_layers: int
resolution: int
sequence_length: int
downsample: str
no_pos_embd: bool
embedding_dim: int = field(default=256)
n_codes: int = field(default=2048)
n_hiddens: int = field(default=240)
n_res_layers: int = field(default=4)
resolution: int = field(default=128)
sequence_length: int = field(default=16)
downsample: Union[str, Tuple[int, int, int]] = field(default=(4,4,4))
no_pos_embd: bool = field(default=True)
def __post_init__(self):
if isinstance(self.downsample, str):
self.downsample = tuple(map(int, self.downsample.split(',')))
+5 -14
View File
@@ -23,18 +23,9 @@ class VQVAETrainingArgument(TrainingArguments):
default=False, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."}
)
def train(args, VQVAE_args, training_args):
def train(args, vqvae_args, training_args):
# Load Model
model_config = VQVAEConfiguration(
embedding_dim=args.embedding_dim,
n_codes=args.n_codes,
n_hiddens=args.n_hiddens,
n_res_layers=args.n_res_layers,
sequence_length=args.sequence_length,
downsample=args.downsample,
resolution=args.resolution
)
model = VQVAEModel(model_config)
model = VQVAEModel(vqvae_args)
# Load Dataset
dataset = VQVAEDataset(args.data_path, sequence_length=args.sequence_length)
@@ -45,8 +36,8 @@ def train(args, VQVAE_args, training_args):
if __name__ == "__main__":
parser = HfArgumentParser((VQVAEArgument, VQVAETrainingArgument))
VQVAE_args, training_args = parser.parse_args_into_dataclasses()
args = argparse.Namespace(**vars(VQVAE_args), **vars(training_args))
vqvae_args, training_args = parser.parse_args_into_dataclasses()
args = argparse.Namespace(**vars(vqvae_args), **vars(training_args))
set_seed(args.seed)
train(args, VQVAE_args, training_args)
train(args, vqvae_args, training_args)