mirror of
https://github.com/facebookresearch/ReAgent.git
synced 2026-06-16 12:44:41 +00:00
Zhengxing Chen
3312c5ce61
Summary: Pull Request resolved: https://github.com/facebookresearch/ReAgent/pull/347 While we have some uninferrable schedulers for FB internal use, we should still have counterparts for OSS. Reviewed By: kaiwenw Differential Revision: D25283564 fbshipit-source-id: 89728d2fb76b5a0278415fc2e5c0fc9dfa0070e7
73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
from reagent.optimizer.uninferrable_optimizers import Adam
|
|
from reagent.optimizer.uninferrable_schedulers import (
|
|
CosineAnnealingLR,
|
|
CosineAnnealingWarmRestarts,
|
|
ExponentialLR,
|
|
MultiStepLR,
|
|
OneCycleLR,
|
|
StepLR,
|
|
)
|
|
from reagent.optimizer.utils import is_torch_lr_scheduler, is_torch_optimizer
|
|
|
|
|
|
class TestMakeOptimizer(unittest.TestCase):
|
|
def setUp(self):
|
|
self.model = torch.nn.Linear(3, 4)
|
|
|
|
def _verify_optimizer(self, optimizer):
|
|
self.assertTrue(is_torch_optimizer(type(optimizer.optimizer)))
|
|
for lr_scheduler in optimizer.lr_schedulers:
|
|
self.assertTrue(is_torch_lr_scheduler(type(lr_scheduler)))
|
|
|
|
def test_make_optimizer_with_step_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(
|
|
lr=0.001, lr_schedulers=[StepLR(gamma=0.1, step_size=0.01)]
|
|
).make_optimizer(self.model.parameters())
|
|
)
|
|
|
|
def test_make_optimizer_with_multistep_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(
|
|
lr=0.001,
|
|
lr_schedulers=[MultiStepLR(gamma=0.2, milestones=[1000, 2000])],
|
|
).make_optimizer(self.model.parameters())
|
|
)
|
|
|
|
def test_make_optimizer_with_exponential_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(lr=0.001, lr_schedulers=[ExponentialLR(gamma=0.9)]).make_optimizer(
|
|
self.model.parameters()
|
|
)
|
|
)
|
|
|
|
def test_make_optimizer_with_cosine_annealing_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(lr=0.001, lr_schedulers=[CosineAnnealingLR(T_max=1)]).make_optimizer(
|
|
self.model.parameters()
|
|
)
|
|
)
|
|
|
|
def test_make_optimizer_with_one_cycle_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(
|
|
lr=0.001,
|
|
lr_schedulers=[
|
|
OneCycleLR(max_lr=0.1, base_momentum=0.8, total_steps=1000)
|
|
],
|
|
).make_optimizer(self.model.parameters())
|
|
)
|
|
|
|
def test_make_optimizer_with_cosine_annealing_warm_restarts_lr_scheduler(self):
|
|
self._verify_optimizer(
|
|
Adam(
|
|
lr=0.001, lr_schedulers=[CosineAnnealingWarmRestarts(T_0=1)]
|
|
).make_optimizer(self.model.parameters())
|
|
)
|