mirror of https://github.com/FoundationVision/GLEE
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
|
|
import unittest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from detectron2.layers import ASPP, DepthwiseSeparableConv2d, FrozenBatchNorm2d
|
|
from detectron2.modeling.backbone.resnet import BasicStem, ResNet
|
|
|
|
|
|
"""
|
|
Test for misc layers.
|
|
"""
|
|
|
|
|
|
class TestBlocks(unittest.TestCase):
|
|
def test_separable_conv(self):
|
|
DepthwiseSeparableConv2d(3, 10, norm1="BN", activation1=nn.PReLU())
|
|
|
|
def test_aspp(self):
|
|
m = ASPP(3, 10, [2, 3, 4], norm="", activation=nn.PReLU())
|
|
self.assertIsNot(m.convs[0].activation.weight, m.convs[1].activation.weight)
|
|
self.assertIsNot(m.convs[0].activation.weight, m.project.activation.weight)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
|
def test_frozen_batchnorm_fp16(self):
|
|
from torch.cuda.amp import autocast
|
|
|
|
C = 10
|
|
input = torch.rand(1, C, 10, 10).cuda()
|
|
m = FrozenBatchNorm2d(C).cuda()
|
|
with autocast():
|
|
output = m(input.half())
|
|
self.assertEqual(output.dtype, torch.float16)
|
|
|
|
# requires_grad triggers a different codepath
|
|
input.requires_grad_()
|
|
with autocast():
|
|
output = m(input.half())
|
|
self.assertEqual(output.dtype, torch.float16)
|
|
|
|
def test_resnet_unused_stages(self):
|
|
resnet = ResNet(BasicStem(), ResNet.make_default_stages(18), out_features=["res2"])
|
|
self.assertTrue(hasattr(resnet, "res2"))
|
|
self.assertFalse(hasattr(resnet, "res3"))
|
|
self.assertFalse(hasattr(resnet, "res5"))
|
|
|
|
resnet = ResNet(BasicStem(), ResNet.make_default_stages(18), out_features=["res2", "res5"])
|
|
self.assertTrue(hasattr(resnet, "res2"))
|
|
self.assertTrue(hasattr(resnet, "res4"))
|
|
self.assertTrue(hasattr(resnet, "res5"))
|