File size: 1,404 Bytes
fb238e8 |
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 |
from pytorch_lightning import (
Trainer,
)
from fengshen.models.model_utils import (
add_module_args,
)
import argparse
from fengshen.data.universal_datamodule import UniversalDataModule
from fengshen.utils.universal_checkpoint import UniversalCheckpoint
from fengshen.examples.pretrain_taiyi_clip.pretrain import (
TaiyiCLIP,
Collator,
)
from fengshen.data.fs_datasets import load_dataset
from torch.utils.data import DataLoader
if __name__ == '__main__':
args_parser = argparse.ArgumentParser()
args_parser = add_module_args(args_parser)
args_parser = UniversalDataModule.add_data_specific_args(args_parser)
args_parser = Trainer.add_argparse_args(args_parser)
args_parser = TaiyiCLIP.add_module_specific_args(args_parser)
args_parser = UniversalCheckpoint.add_argparse_args(args_parser)
args = args_parser.parse_args()
checkpoint_callback = UniversalCheckpoint(args)
trainer = Trainer.from_argparse_args(args, callbacks=[
checkpoint_callback
])
model = TaiyiCLIP(args)
processor = model.processor
collate_fn = Collator(processor)
datasets = load_dataset(args.datasets_name)
dataloader = DataLoader(datasets[args.test_datasets_field],
batch_size=args.test_batchsize, num_workers=2, collate_fn=collate_fn)
trainer.validate(model, dataloaders=dataloader, ckpt_path=args.load_ckpt_path)
|