def train(datasets, model, loss, optimizer, meters, args): task = (ClassificationTask() .set_num_epochs(args.num_epochs) .set_loss(loss) .set_model(model) .set_optimizer(optimizer) .set_meters(meters)) for phase in ["train", "test"]: task.set_dataset(datasets[phase], phase) hooks = [LossLrMeterLoggingHook(log_freq=args.print_freq)] # show progress hooks.append(ProgressBarHook()) if not args.skip_tensorboard: try: from tensorboardX import SummaryWriter tb_writer = SummaryWriter(log_dir=args.video_dir + "/tensorboard") hooks.append(TensorboardPlotHook(tb_writer)) except ImportError: print("tensorboardX not installed, skipping tensorboard hooks") checkpoint_dir = f"{args.video_dir}/checkpoint/classy_checkpoint_{time.time()}" os.mkdir(checkpoint_dir) hooks.append(CheckpointHook(checkpoint_dir, input_args={})) task = task.set_hooks(hooks) trainer = LocalTrainer(use_gpu=args.cuda, num_dataloader_workers=args.num_workers) trainer.train(task)
def test_one(self): train_dataset = MyDataset( batchsize_per_replica=32, shuffle=False, transform=GenericImageTransform( transform=transforms.Compose( [ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ), ] ) ), num_samples=100, crop_size=224, class_ratio=0.5, seed=0, ) test_dataset = MyDataset( batchsize_per_replica=32, shuffle=False, transform=GenericImageTransform( transform=transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ), ] ) ), num_samples=100, crop_size=224, class_ratio=0.5, seed=0, ) model = MyModel() loss = MyLoss() optimizer = SGD(lr_scheduler=ConstantParamScheduler(0.01)) task = ( ClassificationTask() .set_model(model) .set_dataset(train_dataset, "train") .set_dataset(test_dataset, "test") .set_loss(loss) .set_optimizer(optimizer) .set_num_epochs(1) ) trainer = LocalTrainer() trainer.train(task)
def _build_task(self, num_epochs, skip_param_schedulers=False): config = self._get_config(skip_param_schedulers) config["optimizer"]["num_epochs"] = num_epochs task = (ClassificationTask().set_num_epochs(num_epochs).set_loss( build_loss(config["loss"])).set_model(build_model( config["model"])).set_optimizer( build_optimizer(config["optimizer"]))) for phase_type in ["train", "test"]: dataset = build_dataset(config["dataset"][phase_type]) task.set_dataset(dataset, phase_type) self.assertTrue(task is not None) return task
def test_get_state(self): config = get_test_task_config() loss = build_loss(config["loss"]) task = ( ClassificationTask().set_num_epochs(1).set_loss(loss).set_model( build_model(config["model"])).set_optimizer( build_optimizer(config["optimizer"]))) for phase_type in ["train", "test"]: dataset = build_dataset(config["dataset"][phase_type]) task.set_dataset(dataset, phase_type) task.prepare() task = build_task(config) task.prepare()
def test_training(self): """Checks we can train a small MLP model.""" config = get_test_mlp_task_config() task = (ClassificationTask().set_num_epochs(10).set_loss( build_loss(config["loss"])).set_model(build_model( config["model"])).set_optimizer( build_optimizer(config["optimizer"])).set_meters([ AccuracyMeter(topk=[1]) ]).set_hooks([LossLrMeterLoggingHook()])) for split in ["train", "test"]: dataset = build_dataset(config["dataset"][split]) task.set_dataset(dataset, split) self.assertTrue(task is not None) trainer = LocalTrainer() trainer.train(task) accuracy = task.meters[0].value["top_1"] self.assertAlmostEqual(accuracy, 1.0)