xdecoder/Instruct-X-Decoder

1import logging2import os3import time4import pickle5 6import torch7import torch.distributed as dist8 9from fvcore.nn import FlopCountAnalysis10from fvcore.nn import flop_count_table11from fvcore.nn import flop_count_str12 13logger = logging.getLogger(__name__)14 15 16NORM_MODULES = [17    torch.nn.BatchNorm1d,18    torch.nn.BatchNorm2d,19    torch.nn.BatchNorm3d,20    torch.nn.SyncBatchNorm,21    # NaiveSyncBatchNorm inherits from BatchNorm2d22    torch.nn.GroupNorm,23    torch.nn.InstanceNorm1d,24    torch.nn.InstanceNorm2d,25    torch.nn.InstanceNorm3d,26    torch.nn.LayerNorm,27    torch.nn.LocalResponseNorm,28]29 30def register_norm_module(cls):31    NORM_MODULES.append(cls)32    return cls