File size: 1,835 Bytes
d670799 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# Copyright (c) OpenMMLab. All rights reserved.
from unittest.mock import MagicMock
import torch
from mmaction.registry import MODELS
from mmaction.structures import ActionDataSample
from mmaction.testing import get_recognizer_cfg
from mmaction.utils import register_all_modules
def test_omni_resnet():
register_all_modules()
config = get_recognizer_cfg(
'omnisource/slowonly_r50_8xb16-8x8x1-256e_imagenet-kinetics400-rgb.py')
recognizer = MODELS.build(config.model)
# test train_step
video_sample = {
'inputs': [
torch.randint(0, 255, (1, 3, 8, 224, 224)),
torch.randint(0, 255, (1, 3, 8, 224, 224))
],
'data_samples': [
ActionDataSample().set_gt_label(2),
ActionDataSample().set_gt_label(2)
]
}
image_sample = {
'inputs': [
torch.randint(0, 255, (1, 3, 224, 224)),
torch.randint(0, 255, (1, 3, 224, 224))
],
'data_samples': [
ActionDataSample().set_gt_label(2),
ActionDataSample().set_gt_label(2)
]
}
optim_wrapper = MagicMock()
loss_vars = recognizer.train_step([video_sample, image_sample],
optim_wrapper)
assert 'loss_cls_0' in loss_vars
assert 'loss_cls_1' in loss_vars
loss_vars = recognizer.train_step([image_sample, video_sample],
optim_wrapper)
assert 'loss_cls_0' in loss_vars
assert 'loss_cls_1' in loss_vars
# test test_step
with torch.no_grad():
predictions = recognizer.test_step(video_sample)
score = predictions[0].pred_score
assert len(predictions) == 2
assert torch.min(score) >= 0
assert torch.max(score) <= 1
|