<gh_stars>1-10
import os
import json
import multiprocessing
import random
import math
from math import log2, floor
from functools import partial
from contextlib import contextmanager, ExitStack
from pathlib import Path
from shutil import rmtree
from collections import Counter
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import torch
import torchvision
from torch.cuda.amp import autocast, GradScaler
from torch.optim import Adam
from torch import nn, einsum
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.autograd import grad as torch_grad
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import Sampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.dataset import Dataset


from PIL import Image
import pandas as pd
from kornia import filter2D
from tensorboardX import SummaryWriter
import numpy as np
import torch.distributed as dist
from tqdm import tqdm
import torch.multiprocessing as mp
import fire
from einops import rearrange
from gsa_pytorch import GSA
from linear_attention_transformer import ImageLinearAttention
from torchvision import transforms

import os
import json
import multiprocessing
import random
import math
from math import log2, floor
from functools import partial
from contextlib import contextmanager, ExitStack
from pathlib import Path
from shutil import rmtree
from collections import Counter
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import torch
import torchvision
from torch.cuda.amp import autocast, GradScaler
from torch.optim import Adam
from torch import nn, einsum
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.autograd import grad as torch_grad
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import Sampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.dataset import Dataset


from PIL import Image
import pandas as pd
from kornia import filter2D
from tensorboardX import SummaryWriter
import numpy as np
import torch.distributed as dist
from tqdm import tqdm
import torch.multiprocessing as mp
import fire
from einops import rearrange
from gsa_pytorch import GSA
from linear_attention_transformer import ImageLinearAttention
from torchvision import transforms