|
|
import os |
|
|
|
|
|
import hydra |
|
|
import hydra.experimental |
|
|
import numpy as np |
|
|
import pytest |
|
|
import torch |
|
|
|
|
|
pytest_plugins = ["helpers_namespace"] |
|
|
|
|
|
hydra.experimental.initialize( |
|
|
os.path.join(os.path.dirname(__file__), "../pointnet2/config") |
|
|
) |
|
|
|
|
|
|
|
|
@pytest.helpers.register |
|
|
def build_cfg(overrides=[]): |
|
|
return hydra.experimental.compose("config.yaml", overrides) |
|
|
|
|
|
|
|
|
@pytest.helpers.register |
|
|
def get_model(overrides=[]): |
|
|
cfg = build_cfg(overrides) |
|
|
return hydra.utils.instantiate(cfg.task_model, cfg) |
|
|
|
|
|
|
|
|
def _test_loop(model, inputs, labels): |
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5) |
|
|
|
|
|
prev_loss = 1e10 |
|
|
for _ in range(5): |
|
|
optimizer.zero_grad() |
|
|
res = model.training_step((inputs, labels), None) |
|
|
loss = res["loss"] |
|
|
loss.backward() |
|
|
optimizer.step() |
|
|
|
|
|
assert loss.item() < prev_loss + 1.0, "Loss spiked upwards" |
|
|
|
|
|
prev_loss = loss.item() |
|
|
|
|
|
|
|
|
@pytest.helpers.register |
|
|
def cls_test(model): |
|
|
B, N = 4, 2048 |
|
|
inputs = torch.randn(B, N, 6).cuda() |
|
|
labels = torch.from_numpy(np.random.randint(0, 3, size=B)).cuda() |
|
|
model.cuda() |
|
|
|
|
|
_test_loop(model, inputs, labels) |
|
|
|
|
|
|
|
|
@pytest.helpers.register |
|
|
def semseg_test(model): |
|
|
B, N = 4, 2048 |
|
|
inputs = torch.randn(B, N, 9).cuda() |
|
|
labels = torch.from_numpy(np.random.randint(0, 3, size=B * N)).view(B, N).cuda() |
|
|
model.cuda() |
|
|
|
|
|
_test_loop(model, inputs, labels) |
|
|
|
