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 train_with_clipped_gradients(amp_args=None):
task = build_task(get_fast_test_task_config())
task.set_num_epochs(1)
task.set_model(SimpleModel())
task.set_loss(SimpleLoss())
task.set_meters([])
task.set_use_gpu(torch.cuda.is_available())
task.set_clip_grad_norm(0.5)
task.set_amp_args(amp_args)
task.set_optimizer(SGD(lr=1))
trainer = LocalTrainer()
trainer.train(task)
return task.model.param.grad.norm()
def train_with_batch(self, simulated_bs, actual_bs, clip_grad_norm=None):
config = copy.deepcopy(get_fast_test_task_config())
config["dataset"]["train"]["num_samples"] = 12
config["dataset"]["train"]["batchsize_per_replica"] = actual_bs
del config["dataset"]["test"]
task = build_task(config)
task.set_num_epochs(1)
task.set_model(SimpleModel())
task.set_loss(SimpleLoss())
task.set_meters([])
task.set_use_gpu(torch.cuda.is_available())
if simulated_bs is not None:
task.set_simulated_global_batchsize(simulated_bs)
if clip_grad_norm is not None:
task.set_clip_grad_norm(clip_grad_norm)
task.set_optimizer(SGD(lr=1))
trainer = LocalTrainer()
trainer.train(task)
return task.model.param