(cfg: DictConfig)
| 274 | |
| 275 | |
| 276 | def train(cfg: DictConfig): |
| 277 | if cfg.ssl_name == "anytouch": |
| 278 | if cfg.load_from_clip: |
| 279 | cfg.data.dataset.config.num_frames = 2 |
| 280 | cfg.data.dataset.config.frame_stride = 6 |
| 281 | if cfg.two_frame: |
| 282 | cfg.data.dataset.config.num_frames = 2 |
| 283 | cfg.data.dataset.config.frame_stride = 6 |
| 284 | if cfg.input_diff: |
| 285 | cfg.data.dataset.config.remove_bg = True |
| 286 | |
| 287 | resume_state, cfg = attempt_resume(cfg) |
| 288 | |
| 289 | logger.info("Instantiating wandb ...") |
| 290 | wandb = init_wandb(cfg.wandb) |
| 291 | if not resume_state: |
| 292 | wandb.config.update(OmegaConf.to_container(cfg, resolve=True)) |
| 293 | OmegaConf.save(cfg, f"{cfg.paths.output_dir}/config.yaml") |
| 294 | |
| 295 | print_config_tree(cfg, resolve=True, save_to_file=True) |
| 296 | if cfg.get("seed"): |
| 297 | seed_everything(cfg.seed, workers=True) |
| 298 | _GLOBAL_SEED = cfg.seed |
| 299 | np.random.seed(_GLOBAL_SEED) |
| 300 | torch.manual_seed(_GLOBAL_SEED) |
| 301 | torch.backends.cudnn.benchmark = True |
| 302 | |
| 303 | logger.info( |
| 304 | f"Instantiating dataset & dataloaders for <{cfg.data.dataset._target_}>" |
| 305 | ) |
| 306 | train_dataloader, val_dataloader = get_dataloaders(cfg) |
| 307 | |
| 308 | trainer = Trainer(wandb_logger=wandb, **cfg.trainer) |
| 309 | |
| 310 | logger.info(f"Instantiating model <{cfg.task._target_}>") |
| 311 | |
| 312 | if cfg.ssl_name == "anytouch": |
| 313 | from transformers import AutoConfig |
| 314 | |
| 315 | num_frames = cfg.data.dataset.config.num_frames |
| 316 | frame_stride = cfg.data.dataset.config.frame_stride |
| 317 | |
| 318 | if cfg.size == 'base': |
| 319 | clip_config_path = os.path.join(os.path.dirname(__file__), '..', 'CLIP-B-16') |
| 320 | config = AutoConfig.from_pretrained(clip_config_path) |
| 321 | else: |
| 322 | raise ValueError(f"Unknown size {cfg.size} for AnyTouch model") |
| 323 | |
| 324 | mae_args = argparse.Namespace(mask_ratio=0.0, stride=frame_stride) |
| 325 | base_encoder = TactileVideoMAE(mae_args, config, num_frames, tube_size=1) |
| 326 | base_encoder = load_model_from_multi_clip( |
| 327 | torch.load(cfg.ckpt_path, map_location='cpu'), base_encoder |
| 328 | ) |
| 329 | print(f"Loaded AnyTouch model from {cfg.ckpt_path} with size {cfg.size}") |
| 330 | |
| 331 | sensor_int = _SENSOR_TYPE_MAP.get(cfg.data.sensor, 1) |
| 332 | model_encoder = _AnyTouchEncoderWrapper(base_encoder, sensor_int) |
| 333 |
no test coverage detected