|
|
| import torch
|
|
|
| from mmaction.models import OmniHead
|
|
|
|
|
| class obj():
|
|
|
| def __init__(self, name, value):
|
| super(obj, self).__init__()
|
| setattr(self, name, value)
|
|
|
|
|
| def testOmniHead():
|
| head = OmniHead(image_classes=100, video_classes=200, in_channels=400)
|
|
|
| image_feat = torch.randn(2, 400, 8, 8)
|
| image_score = head(image_feat)
|
| assert image_score.shape == torch.Size([2, 100])
|
|
|
| video_feat = torch.randn(2, 400, 8, 8, 8)
|
| video_score = head(video_feat)
|
| assert video_score.shape == torch.Size([2, 200])
|
|
|
| head = OmniHead(
|
| image_classes=100,
|
| video_classes=200,
|
| in_channels=400,
|
| video_nl_head=True)
|
|
|
| video_feat = torch.randn(2, 400, 8, 8, 8)
|
| video_score = head(video_feat)
|
| assert video_score.shape == torch.Size([2, 200])
|
| data_samples = [obj('gt_label', torch.tensor(1)) for _ in range(2)]
|
| losses = head.loss_by_feat(video_score, data_samples)
|
| assert 'loss_cls' in losses
|
|
|
| image_feat = torch.randn(1, 400, 8, 8)
|
| head.eval()
|
| image_score = head(image_feat)
|
| assert image_score.shape == torch.Size([1, 100])
|
| data_samples = [obj('gt_label', torch.tensor(1))]
|
| losses = head.loss_by_feat(image_score, data_samples)
|
| assert 'loss_cls' in losses
|
|
|
