[refactor] configuration default values
Former-commit-id: d7ca7685d6a8c6d015082dda241bc98ab156770d
This commit is contained in:
@@ -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(',')))
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user