codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch nat model import collections import unittest from transformers import natconfig from transformers testingutils import requirenatten requiretorch requirevision slow torchdevice from transformers utils import cachedproperty istorchavailable isvisionavailable from testbackbonecommon import backbonetestermixin from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin configzeroinit floatstensor idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from torch import nn from transformers import natbackbone natforimageclassification natmodel from transformers models nat modelingnat import natpretrainedmodelarchivelist if isvisionavailable from pil import image from transformers import autoimageprocessor class natmodeltester def init self parent batchsize13 imagesize64 patchsize4 numchannels3 embeddim16 depths1 2 1 numheads2 4 8 kernelsize3 mlpratio2 0 qkvbiastrue hiddendropoutprob0 0 attentionprobsdropoutprob0 0 droppathrate0 1 hiddenactgelu patchnormtrue initializerrange0 02 layernormeps1e5 istrainingtrue scopenone uselabelstrue numlabels10 outfeaturesstage1 stage2 outindices1 2 self parent parent self batchsize batchsize self imagesize imagesize self patchsize patchsize self numchannels numchannels self embeddim embeddim self depths depths self numheads numheads self kernelsize kernelsize self mlpratio mlpratio self qkvbias qkvbias self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self droppathrate droppathrate self hiddenact hiddenact self patchnorm patchnorm self layernormeps layernormeps self initializerrange initializerrange self istraining istraining self scope scope self uselabels uselabels self numlabels numlabels self outfeatures outfeatures self outindices outindices def prepareconfigandinputsself pixelvalues floatstensorself batchsize self numchannels self imagesize self imagesize labels none if self uselabels labels idstensorself batchsize self numlabels config self getconfig return config pixelvalues labels def getconfigself return natconfig numlabelsself numlabels imagesizeself imagesize patchsizeself patchsize numchannelsself numchannels embeddimself embeddim depthsself depths numheadsself numheads kernelsizeself kernelsize mlpratioself mlpratio qkvbiasself qkvbias hiddendropoutprobself hiddendropoutprob attentionprobsdropoutprobself attentionprobsdropoutprob droppathrateself droppathrate hiddenactself hiddenact patchnormself patchnorm layernormepsself layernormeps initializerrangeself initializerrange outfeaturesself outfeatures outindicesself outindices def createandcheckmodelself config pixelvalues labels model natmodelconfigconfig model totorchdevice model eval result modelpixelvalues expectedheight expectedwidth config imagesize config patchsize 2 lenconfig depths 1 expecteddim intconfig embeddim 2 lenconfig depths 1 self parent assertequal result lasthiddenstate shape self batchsize expectedheight expectedwidth expecteddim def createandcheckforimageclassificationself config pixelvalues labels model natforimageclassificationconfig model totorchdevice model eval result modelpixelvalues labelslabels self parent assertequalresult logits shape self batchsize self numlabels test greyscale images config numchannels 1 model natforimageclassificationconfig model totorchdevice model eval pixelvalues floatstensorself batchsize 1 self imagesize self imagesize result modelpixelvalues self parent assertequalresult logits shape self batchsize self numlabels def createandcheckbackboneself config pixelvalues labels model natbackboneconfigconfig model totorchdevice model eval result modelpixelvalues verify hidden states self parent assertequallenresult featuremaps lenconfig outfeatures self parent assertlistequallistresult featuremaps0 shape self batchsize model channels0 16 16 verify channels self parent assertequallenmodel channels lenconfig outfeatures verify backbone works with outfeaturesnone config outfeatures none model natbackboneconfigconfig model totorchdevice model eval result modelpixelvalues verify feature maps self parent assertequallenresult featuremaps 1 self parent assertlistequallistresult featuremaps0 shape self batchsize model channels1 4 4 verify channels self parent assertequallenmodel channels 1 def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config pixelvalues labels configandinputs inputsdict pixelvalues pixelvalues return config inputsdict requirenatten requiretorch class natmodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses natmodel natforimageclassification natbackbone if istorchavailable else pipelinemodelmapping featureextraction natmodel imageclassification natforimageclassification if istorchavailable else fxcompatible false testtorchscript false testpruning false testresizeembeddings false testheadmasking false def setupself self modeltester natmodeltesterself self configtester configtesterself configclassnatconfig embeddim37 def testconfigself self createandtestconfigcommonproperties self configtester createandtestconfigtojsonstring self configtester createandtestconfigtojsonfile self configtester createandtestconfigfromandsavepretrained self configtester createandtestconfigwithnumlabels self configtester checkconfigcanbeinitwithoutparams self configtester checkconfigargumentsinit def createandtestconfigcommonpropertiesself return def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testforimageclassificationself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforimageclassificationconfigandinputs def testbackboneself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckbackboneconfigandinputs unittest skipreasonnat does not use inputsembeds def testinputsembedsself pass unittest skipreasonnat does not use feedforward chunking def testfeedforwardchunkingself pass def testmodelcommonattributesself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig self assertisinstancemodel getinputembeddings nn module x model getoutputembeddings self asserttruex is none or isinstancex nn linear def testattentionoutputsself self skiptestnat s attention operation is handled entirely by natten def checkhiddenstatesoutputself inputsdict config modelclass imagesize model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass hiddenstates outputs hiddenstates expectednumlayers getattr self modeltester expectednumhiddenlayers lenself modeltester depths 1 self assertequallenhiddenstates expectednumlayers nat has a different seqlength patchsize config patchsize if isinstanceconfig patchsize collections abc iterable else config patchsize config patchsize height imagesize0 patchsize0 width imagesize1 patchsize1 self assertlistequal listhiddenstates0 shape3 height width self modeltester embeddim if modelclass name natbackbone reshapedhiddenstates outputs reshapedhiddenstates self assertequallenreshapedhiddenstates expectednumlayers batchsize numchannels height width reshapedhiddenstates0 shape reshapedhiddenstates reshapedhiddenstates0 viewbatchsize numchannels height width permute0 2 3 1 self assertlistequal listreshapedhiddenstates shape3 height width self modeltester embeddim def testhiddenstatesoutputself config inputsdict self modeltester prepareconfigandinputsforcommon imagesize self modeltester imagesize if isinstanceself modeltester imagesize collections abc iterable else self modeltester imagesize self modeltester imagesize for modelclass in self allmodelclasses inputsdictoutputhiddenstates true self checkhiddenstatesoutputinputsdict config modelclass imagesize check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true self checkhiddenstatesoutputinputsdict config modelclass imagesize slow def testmodelfrompretrainedself for modelname in natpretrainedmodelarchivelist 1 model natmodel frompretrainedmodelname self assertisnotnonemodel def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters if embeddings not in name and param requiresgrad self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized requirenatten requirevision requiretorch class natmodelintegrationtestunittest testcase cachedproperty def defaultimageprocessorself return autoimageprocessor frompretrainedshilabsnatminiin1k224 if isvisionavailable else none slow def testinferenceimageclassificationheadself model natforimageclassification frompretrainedshilabsnatminiin1k224 totorchdevice imageprocessor self defaultimageprocessor image image open testsfixturestestssamplescoco000000039769 png inputs imageprocessorimagesimage returntensorspt totorchdevice forward pass with torch nograd outputs modelinputs verify the logits expectedshape torch size1 1000 self assertequaloutputs logits shape expectedshape expectedslice torch tensor0 3805 0 8676 0 3912 totorchdevice self asserttruetorch allcloseoutputs logits0 3 expectedslice atol1e4 requiretorch requirenatten class natbackbonetestunittest testcase backbonetestermixin allmodelclasses natbackbone if istorchavailable else configclass natconfig def setupself self modeltester natmodeltesterself coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch nat model test greyscale images verify hidden states verify channels verify backbone works with out_features none verify feature maps verify channels nat has a different seq_length check that output_hidden_states also work using config forward pass verify the logits
import collections
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import NatBackbone, NatForImageClassification, NatModel
    from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor


class NatModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        image_size=64,
        patch_size=4,
        num_channels=3,
        embed_dim=16,
        depths=[1, 2, 1],
        num_heads=[2, 4, 8],
        kernel_size=3,
        mlp_ratio=2.0,
        qkv_bias=True,
        hidden_dropout_prob=0.0,
        attention_probs_dropout_prob=0.0,
        drop_path_rate=0.1,
        hidden_act="gelu",
        patch_norm=True,
        initializer_range=0.02,
        layer_norm_eps=1e-5,
        is_training=True,
        scope=None,
        use_labels=True,
        num_labels=10,
        out_features=["stage1", "stage2"],
        out_indices=[1, 2],
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_channels = num_channels
        self.embed_dim = embed_dim
        self.depths = depths
        self.num_heads = num_heads
        self.kernel_size = kernel_size
        self.mlp_ratio = mlp_ratio
        self.qkv_bias = qkv_bias
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.drop_path_rate = drop_path_rate
        self.hidden_act = hidden_act
        self.patch_norm = patch_norm
        self.layer_norm_eps = layer_norm_eps
        self.initializer_range = initializer_range
        self.is_training = is_training
        self.scope = scope
        self.use_labels = use_labels
        self.num_labels = num_labels
        self.out_features = out_features
        self.out_indices = out_indices

    def prepare_config_and_inputs(self):
        pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

        labels = None
        if self.use_labels:
            labels = ids_tensor([self.batch_size], self.num_labels)

        config = self.get_config()

        return config, pixel_values, labels

    def get_config(self):
        return NatConfig(
            num_labels=self.num_labels,
            image_size=self.image_size,
            patch_size=self.patch_size,
            num_channels=self.num_channels,
            embed_dim=self.embed_dim,
            depths=self.depths,
            num_heads=self.num_heads,
            kernel_size=self.kernel_size,
            mlp_ratio=self.mlp_ratio,
            qkv_bias=self.qkv_bias,
            hidden_dropout_prob=self.hidden_dropout_prob,
            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
            drop_path_rate=self.drop_path_rate,
            hidden_act=self.hidden_act,
            patch_norm=self.patch_norm,
            layer_norm_eps=self.layer_norm_eps,
            initializer_range=self.initializer_range,
            out_features=self.out_features,
            out_indices=self.out_indices,
        )

    def create_and_check_model(self, config, pixel_values, labels):
        model = NatModel(config=config)
        model.to(torch_device)
        model.eval()
        result = model(pixel_values)

        expected_height = expected_width = (config.image_size // config.patch_size) // (2 ** (len(config.depths) - 1))
        expected_dim = int(config.embed_dim * 2 ** (len(config.depths) - 1))

        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, expected_height, expected_width, expected_dim)
        )

    def create_and_check_for_image_classification(self, config, pixel_values, labels):
        model = NatForImageClassification(config)
        model.to(torch_device)
        model.eval()
        result = model(pixel_values, labels=labels)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))

        
        config.num_channels = 1
        model = NatForImageClassification(config)
        model.to(torch_device)
        model.eval()

        pixel_values = floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
        result = model(pixel_values)
        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))

    def create_and_check_backbone(self, config, pixel_values, labels):
        model = NatBackbone(config=config)
        model.to(torch_device)
        model.eval()
        result = model(pixel_values)

        
        self.parent.assertEqual(len(result.feature_maps), len(config.out_features))
        self.parent.assertListEqual(list(result.feature_maps[0].shape), [self.batch_size, model.channels[0], 16, 16])

        
        self.parent.assertEqual(len(model.channels), len(config.out_features))

        
        config.out_features = None
        model = NatBackbone(config=config)
        model.to(torch_device)
        model.eval()
        result = model(pixel_values)

        
        self.parent.assertEqual(len(result.feature_maps), 1)
        self.parent.assertListEqual(list(result.feature_maps[0].shape), [self.batch_size, model.channels[-1], 4, 4])

        
        self.parent.assertEqual(len(model.channels), 1)

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, pixel_values, labels = config_and_inputs
        inputs_dict = {"pixel_values": pixel_values}
        return config, inputs_dict


@require_natten
@require_torch
class NatModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (
            NatModel,
            NatForImageClassification,
            NatBackbone,
        )
        if is_torch_available()
        else ()
    )
    pipeline_model_mapping = (
        {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
        if is_torch_available()
        else {}
    )
    fx_compatible = False

    test_torchscript = False
    test_pruning = False
    test_resize_embeddings = False
    test_head_masking = False

    def setUp(self):
        self.model_tester = NatModelTester(self)
        self.config_tester = ConfigTester(self, config_class=NatConfig, embed_dim=37)

    def test_config(self):
        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()

    def create_and_test_config_common_properties(self):
        return

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_for_image_classification(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*config_and_inputs)

    def test_backbone(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*config_and_inputs)

    @unittest.skip(reason="Nat does not use inputs_embeds")
    def test_inputs_embeds(self):
        pass

    @unittest.skip(reason="Nat does not use feedforward chunking")
    def test_feed_forward_chunking(self):
        pass

    def test_model_common_attributes(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertIsInstance(model.get_input_embeddings(), (nn.Module))
            x = model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(x, nn.Linear))

    def test_attention_outputs(self):
        self.skipTest("Nat's attention operation is handled entirely by NATTEN.")

    def check_hidden_states_output(self, inputs_dict, config, model_class, image_size):
        model = model_class(config)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            outputs = model(**self._prepare_for_class(inputs_dict, model_class))

        hidden_states = outputs.hidden_states

        expected_num_layers = getattr(
            self.model_tester, "expected_num_hidden_layers", len(self.model_tester.depths) + 1
        )
        self.assertEqual(len(hidden_states), expected_num_layers)

        
        patch_size = (
            config.patch_size
            if isinstance(config.patch_size, collections.abc.Iterable)
            else (config.patch_size, config.patch_size)
        )

        height = image_size[0] // patch_size[0]
        width = image_size[1] // patch_size[1]

        self.assertListEqual(
            list(hidden_states[0].shape[-3:]),
            [height, width, self.model_tester.embed_dim],
        )

        if model_class.__name__ != "NatBackbone":
            reshaped_hidden_states = outputs.reshaped_hidden_states
            self.assertEqual(len(reshaped_hidden_states), expected_num_layers)

            batch_size, num_channels, height, width = reshaped_hidden_states[0].shape
            reshaped_hidden_states = (
                reshaped_hidden_states[0].view(batch_size, num_channels, height, width).permute(0, 2, 3, 1)
            )
            self.assertListEqual(
                list(reshaped_hidden_states.shape[-3:]),
                [height, width, self.model_tester.embed_dim],
            )

    def test_hidden_states_output(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        image_size = (
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size, collections.abc.Iterable)
            else (self.model_tester.image_size, self.model_tester.image_size)
        )

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            self.check_hidden_states_output(inputs_dict, config, model_class, image_size)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            self.check_hidden_states_output(inputs_dict, config, model_class, image_size)

    @slow
    def test_model_from_pretrained(self):
        for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = NatModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                if "embeddings" not in name and param.requires_grad:
                    self.assertIn(
                        ((param.data.mean() * 1e9).round() / 1e9).item(),
                        [0.0, 1.0],
                        msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                    )


@require_natten
@require_vision
@require_torch
class NatModelIntegrationTest(unittest.TestCase):
    @cached_property
    def default_image_processor(self):
        return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224") if is_vision_available() else None

    @slow
    def test_inference_image_classification_head(self):
        model = NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224").to(torch_device)
        image_processor = self.default_image_processor

        image = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
        inputs = image_processor(images=image, return_tensors="pt").to(torch_device)

        
        with torch.no_grad():
            outputs = model(**inputs)

        
        expected_shape = torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape, expected_shape)
        expected_slice = torch.tensor([0.3805, -0.8676, -0.3912]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))


@require_torch
@require_natten
class NatBackboneTest(unittest.TestCase, BackboneTesterMixin):
    all_model_classes = (NatBackbone,) if is_torch_available() else ()
    config_class = NatConfig

    def setUp(self):
        self.model_tester = NatModelTester(self)