Official implementation for CIGN
import av
# import torchaudio
import numpy as np
from fractions import Fraction


# def load_audio_torchaudio(fn):
#     data, sr = torchaudio.load(fn)
#     return data, sr


def open_audio_av(path):
    container = av.open(path)
    for stream in container.streams.video:
        stream.codec_context.thread_type = av.codec.context.ThreadType.NONE
        stream.codec_context.thread_count = 1
    for stream in container.streams.audio:
        stream.codec_context.thread_type = av.codec.context.ThreadType.NONE
        stream.codec_context.thread_count = 1
    return container


def load_audio_av(path=None, container=None, rate=None, start_time=None, duration=None, layout="mono"):
    if container is None:
        container = av.open(path)
    audio_stream = container.streams.audio[0]

    # Parse metadata
    _ss = audio_stream.start_time * audio_stream.time_base if audio_stream.start_time is not None else 0.
    _dur = audio_stream.duration * audio_stream.time_base
    _ff = _ss + _dur
    _rate = audio_stream.rate

    if rate is None:
        rate = _rate
    if start_time is None:
        start_time = _ss
    if duration is None:
        duration = _ff - start_time
    duration = min(duration, _ff - start_time)
    end_time = start_time + duration

    resampler = av.audio.resampler.AudioResampler(format="s16p", layout=layout, rate=rate)

    # Read data
    chunks = []
    container.seek(int(start_time * av.time_base))
    for frame in container.decode(audio=0):
        chunk_start_time = frame.pts * frame.time_base
        chunk_end_time = chunk_start_time + Fraction(frame.samples, frame.rate)
        if chunk_end_time < start_time:   # Skip until start time
            continue
        if chunk_start_time > end_time:       # Exit if clip has been extracted
            break

        try:
            frame.pts = None
            if resampler is not None:
                chunks.append((chunk_start_time, resampler.resample(frame).to_ndarray()))
            else:
                chunks.append((chunk_start_time, frame.to_ndarray()))
        except AttributeError:
            break

    # Trim for frame accuracy
    audio = np.concatenate([af[1] for af in chunks], 1)
    ss = int((start_time - chunks[0][0]) * rate)
    t = int(duration * rate)
    if ss < 0:
        audio = np.pad(audio, ((0, 0), (-ss, 0)), 'constant', constant_values=0)
        ss = 0
    audio = audio[:, ss: ss+t]

    # Normalize to [-1, 1]
    audio = audio / np.iinfo(audio.dtype).max

    return audio, rate


def audio_info_av(inpt, audio=None, format=None):
    container = inpt
    if isinstance(inpt, str):
        try:
            container = av.open(inpt, format=format)
        except av.AVError:
            return None, None

    audio_stream = container.streams.audio[audio]
    time_base = audio_stream.time_base
    duration = audio_stream.duration * time_base
    start_time = audio_stream.start_time * time_base
    channels = audio_stream.channels
    fps = audio_stream.rate
    chunk_size = audio_stream.frame_size
    chunks = audio_stream.frames
    meta = {'channels': channels,
            'fps': fps,
            'start_time': start_time,
            'duration': duration,
            'chunks': chunks,
            'chunk_size': chunk_size}
    return meta

audio_io.py
stoneMo/CIGN
import os
import csv
import cv2
import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
from scipy import signal
import random
import json
import pickle
import xml.etree.ElementTree as ET
from audio_io import load_audio_av, open_audio_av
import torch

from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


def load_image(path):
    return Image.open(path).convert('RGB')


def load_waveform(path, dur=3.):
    # Load audio
    audio_ctr = open_audio_av(path)
    audio_dur = audio_ctr.streams.audio[0].duration * audio_ctr.streams.audio[0].time_base
    audio_ss = max(float(audio_dur)/2 - dur/2, 0)
    audio, samplerate = load_audio_av(container=audio_ctr, start_time=audio_ss, duration=dur)

    # To Mono
    audio = np.clip(audio, -1., 1.).mean(0)

    # Repeat if audio is too short
    if audio.shape[0] < samplerate * dur:
        n = int(samplerate * dur / audio.shape[0]) + 1
        audio = np.tile(audio, n)
    waveform = audio[:int(samplerate * dur)]

    return waveform, samplerate

def log_mel_spectrogram(waveform, samplerate):
    frequencies, times, spectrogram = signal.spectrogram(waveform, samplerate, nperseg=512, noverlap=274)
    spectrogram = np.log(spectrogram + 1e-7)
    return spectrogram

def load_all_bboxes(annotation_dir, format='flickr'):
    gt_bboxes = {}
    if format == 'flickr':
        anno_files = os.listdir(annotation_dir)
        for filename in anno_files:
            file = filename.split('.')[0]
            gt = ET.parse(f"{annotation_dir}/{filename}").getroot()
            bboxes = []
            for child in gt:
                for childs in child:
                    bbox = []
                    if childs.tag == 'bbox':
                        for index, ch in enumerate(childs):
                            if index == 0:
                                continue
                            bbox.append(int(224 * int(ch.text)/256))
                    bboxes.append(bbox)
            gt_bboxes[file] = bboxes

    elif format in {'vggss', 'vggsound_single'}:
        with open('metadata/vggss.json') as json_file:
            annotations = json.load(json_file)
        for annotation in annotations:
            bboxes = [(np.clip(np.array(bbox), 0, 1) * 224).astype(int) for bbox in annotation['bbox']]
            gt_bboxes[annotation['file']] = bboxes

    elif format == 'vggsound_duet':
        gt_bboxes_raw = {}
        with open('metadata/vggss.json') as json_file:
            annotations = json.load(json_file)
        for annotation in annotations:
            gt_bboxes_raw[annotation['file']] = annotation['bbox']
        # fns2cls = {item[0]:item[1] for item in csv.reader(open('metadata/vggsound_duet_test.csv'))}
        fns2mix = {item[0]:item[2] for item in csv.reader(open('metadata/vggsound_duet_test.csv'))}
        for annotation in annotations:
            fn = annotation['file']
            fn_mix = fns2mix[fn]
            bboxes = [(np.clip(np.array(bbox), 0, 1) * 224).astype(int) for bbox in gt_bboxes_raw[fn]]
            bboxes_mix = [(np.clip(np.array(bbox_mix), 0, 1) * 224).astype(int) for bbox_mix in gt_bboxes_raw[fn_mix]]
            bboxes_src = [bboxes, bboxes_mix]
            # classes_src = [fns2cls[fn], fns2cls[fn_mix]]
            gt_bboxes[fn] = bboxes_src

    elif format in {'vgginstruments', 'vgginstruments_multi'}:
        anno_files = os.listdir(annotation_dir)
        for filename in anno_files:
            file = filename.split('.')[0]
            gt_bboxes[file] = f"{annotation_dir}/{filename}"

    elif format == 'music_solo':
        with open('metadata/music_solo.json') as json_file:
            annotations = json.load(json_file)
        for annotation in annotations:
            bboxes = [(np.clip(np.array(annotation['bbox']), 0, 1) * 224).astype(int)]
            gt_bboxes[annotation['file']] = bboxes

    elif format == 'music_duet':
        with open('metadata/music_duet.json') as json_file:
            annotations = json.load(json_file)
        for annotation in annotations:
            bboxes_src = [annotation['bbox_src1'], annotation['bbox_src2']]
            classes_src = [annotation['class_src1'], annotation['class_src2']]
            bboxes = [(np.clip(np.array(bbox), 0, 1) * 224).astype(int) for bbox in bboxes_src]
            gt_bboxes[annotation['file']] = [bboxes, classes_src]

    return gt_bboxes


def bbox2gtmap(bboxes, format='flickr'):
    gt_map = np.zeros([224, 224])
    for xmin, ymin, xmax, ymax in bboxes:
        temp = np.zeros([224, 224])
        temp[ymin:ymax, xmin:xmax] = 1
        gt_map += temp

    if format == 'flickr':
        # Annotation consensus
        gt_map = gt_map / 2
        gt_map[gt_map > 1] = 1

    elif format in {'vggss', 'music_duet'}:
        # Single annotation
        gt_map[gt_map > 0] = 1

    return gt_map


def mask2gtmap(gt_mask_path):
    with open(gt_mask_path, 'rb') as f:
        gt_mask = pickle.load(f)
    gt_map = cv2.resize(gt_mask, (224,224), interpolation=cv2.INTER_NEAREST)
    return gt_map


class AudioVisualDataset(Dataset):
    def __init__(self, image_files, audio_files, image_path, audio_path, mode='train', 
            sup_image_path=None, sup_audio_path=None, audio_dur=3., 
            image_transform=None, audio_transform=None, all_bboxes=None, bbox_format='flickr', 
            num_classes=0, class_labels=None, num_mixtures=1, class_labels_ss=None, 
            image_files_ss=None, audio_files_ss=None, all_bboxes_ss=None):
        super().__init__()
        self.audio_path = audio_path
        self.image_path = image_path

        self.mode = mode
        self.sup_audio_path = sup_audio_path
        self.sup_image_path = sup_image_path

        self.audio_dur = audio_dur

        self.audio_files = audio_files
        self.image_files = image_files
        self.all_bboxes = all_bboxes
        self.bbox_format = bbox_format
        self.class_labels = class_labels
        self.num_classes = num_classes

        self.num_mixtures = num_mixtures
        self.class_labels_ss = class_labels_ss
        self.image_files_ss = image_files_ss
        self.audio_files_ss = audio_files_ss
        self.all_bboxes_ss = all_bboxes_ss

        self.image_transform = image_transform
        self.audio_transform = audio_transform

    def getitem(self, idx):

        image_path = self.image_path
        audio_path = self.audio_path

        anno = {}

        if self.class_labels is not None:
            class_label = torch.zeros(self.num_classes)
            class_idx = self.class_labels[idx]
            class_label[class_idx] = 1
            anno['class'] = class_label

            if self.class_labels_ss is not None:
                class_label_mix = torch.zeros(self.num_classes)
                class_idx_mix = self.class_labels_ss[idx]
                class_label_mix[class_idx_mix] = 1
                anno['class'] = torch.stack([class_label, class_label_mix])

        file = self.image_files[idx]
        file_id = file.split('.')[0]

        # Image
        img_fn = image_path + self.image_files[idx]
        frame = self.image_transform(load_image(img_fn))

        # Audio
        audio_fn = audio_path + self.audio_files[idx]
        waveform, samplerate = load_waveform(audio_fn)
        spectrogram = self.audio_transform(log_mel_spectrogram(waveform, samplerate))

        return frame, spectrogram, anno, file_id

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        try:
            return self.getitem(idx)
        except Exception:
            return self.getitem(random.sample(range(len(self)), 1)[0])


def get_train_dataset(args):
    audio_path = f"{args.train_data_path}/audio/"
    image_path = f"{args.train_data_path}/frames/"

    # List directory
    audio_files = {fn.split('.wav')[0] for fn in os.listdir(audio_path) if fn.endswith('.wav')}
    image_files = {fn.split('.jpg')[0] for fn in os.listdir(image_path) if fn.endswith('.jpg')}
    if args.trainset in {'music_solo', 'music_duet'}:
        avail_audio_files = []
        for image_file in image_files:
            if image_file[:-10] in audio_files:
                avail_audio_files.append(image_file)
        audio_files = {file for file in avail_audio_files}
    avail_files = audio_files.intersection(image_files)
    print(f"{len(avail_files)} available files")

    # Subsample if specified
    trainset_name_list = args.trainset.split(',')
    trainset = []
    for trainset_name in trainset_name_list:
        if trainset_name in {'vgginstruments_group_0', 'vgginstruments_group_1', 'vgginstruments_group_2', 'vgginstruments_group_3'}:
            trainset_list = [line.split(',')[0] for line in open(f"metadata/avctl_{trainset_name}_train.txt").read().splitlines()]
        else:
            trainset_list = open(f"metadata/{args.trainset}.txt").read().splitlines()
        trainset.extend(trainset_list)
    trainset = set(trainset)
    avail_files = avail_files.intersection(trainset)
    print(f"{len(avail_files)} valid subset files")
    
    avail_files = sorted(list(avail_files))
    audio_files = [dt+'.wav' for dt in avail_files]
    image_files = [dt+'.jpg' for dt in avail_files]
    
    all_bboxes = [[] for _ in range(len(image_files))]

    all_classes = ['flute', 'cello', 'bass_guitar', 'accordion', 'tabla', 'erhu', 'cornet', 'electronic_organ', 'timbales', 
    'acoustic_guitar', 'violin', 'piano', 'banjo', 'glockenspiel', 'steel_guitar', 'vibraphone', 'trumpet', 'zither', 
    'cymbal', 'xylophone', 'harp', 'hammond_organ', 'harpsichord', 'bongo', 'bass_drum', 'mandolin', 'guiro', 
    'saxophone', 'electric_guitar', 'drum_kit', 'clarinet', 'snare_drum', 'ukulele', 'sitar', 'double_bass', 'congas']

    fns2cls = {}
    for trainset_name in trainset_name_list:
        if trainset_name in {'vgginstruments_group_0', 'vgginstruments_group_1', 'vgginstruments_group_2', 'vgginstruments_group_3'}:
            fns2cls.update({line.split(',')[0]:line.split(',')[1] for line in open(f"metadata/avctl_{trainset_name}_train.txt").read().splitlines()})
            
    class_labels = []
    for dt in avail_files:
        cls = all_classes.index(fns2cls[dt])
        class_labels.append(cls)

    num_classes = args.num_class

    print('class_labels:', class_labels[:10])
    print('all_classes:', all_classes)

    # Transforms
    image_transform = transforms.Compose([
        transforms.Resize(int(224 * 1.1), Image.BICUBIC),
        transforms.RandomCrop((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])])
    audio_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.0], std=[12.0])])

    return AudioVisualDataset(
        mode='train',
        image_files=image_files,
        audio_files=audio_files,
        all_bboxes=all_bboxes,
        image_path=image_path,
        audio_path=audio_path,
        audio_dur=3.,
        image_transform=image_transform,
        audio_transform=audio_transform,
        num_classes=num_classes,
        class_labels=class_labels,
        num_mixtures=1
    )


def get_test_dataset(args, mode):
    audio_path = args.test_data_path + 'audio/'
    image_path = args.test_data_path + 'frames/'

    #  Retrieve list of audio and video files
    testset_name_list = args.testset.split(',')
    testset = []
    for testset_name in testset_name_list:
        if testset_name in {'vgginstruments_group_0', 'vgginstruments_group_1', 'vgginstruments_group_2', 'vgginstruments_group_3'}:
            testset_list = [line.split(',')[0] for line in open(f"metadata/avctl_{testset_name}_{mode}.txt").read().splitlines()]
            testset.extend(testset_list)
    testset = set(testset)
    
    # Intersect with available files
    audio_files = {fn.split('.wav')[0] for fn in os.listdir(audio_path)}
    image_files = {fn.split('.jpg')[0] for fn in os.listdir(image_path)}
    if args.testset in {'music_solo', 'music_duet'}:
        avail_audio_files = []
        for image_file in image_files:
            if image_file[:-10] in audio_files:
                avail_audio_files.append(image_file)
        audio_files = {file for file in avail_audio_files}

    avail_files = audio_files.intersection(image_files)
    testset = testset.intersection(avail_files)
    print(f"{len(testset)} files for testing")

    testset = sorted(list(testset))
    image_files = [dt+'.jpg' for dt in testset]
    if args.testset in {'music_solo', 'music_duet'}:
        audio_files = [dt[:-10]+'.wav' for dt in testset]
    else:
        audio_files = [dt+'.wav' for dt in testset]
    
    all_classes = ['flute', 'cello', 'bass_guitar', 'accordion', 'tabla', 'erhu', 'cornet', 'electronic_organ', 'timbales', 
    'acoustic_guitar', 'violin', 'piano', 'banjo', 'glockenspiel', 'steel_guitar', 'vibraphone', 'trumpet', 'zither', 
    'cymbal', 'xylophone', 'harp', 'hammond_organ', 'harpsichord', 'bongo', 'bass_drum', 'mandolin', 'guiro', 
    'saxophone', 'electric_guitar', 'drum_kit', 'clarinet', 'snare_drum', 'ukulele', 'sitar', 'double_bass', 'congas']

    fns2cls = {}
    for testset_name in testset_name_list:
        if testset_name in {'vgginstruments_group_0', 'vgginstruments_group_1', 'vgginstruments_group_2', 'vgginstruments_group_3'}:
            fns2cls.update({line.split(',')[0]:line.split(',')[1] for line in open(f"metadata/avctl_{testset_name}_{mode}.txt").read().splitlines()})
            
    class_labels = []
    for dt in testset:
        cls = all_classes.index(fns2cls[dt])
        class_labels.append(cls)

    num_classes = len(list(set(class_labels)))
        
    # Transforms
    image_transform = transforms.Compose([
        transforms.Resize((224, 224), Image.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])])
    audio_transform = transforms.Compose([transforms.ToTensor(),
                                          transforms.Normalize(mean=[0.0], std=[12.0])])

    return AudioVisualDataset(
        mode='test',
        image_files=image_files,
        audio_files=audio_files,
        image_path=image_path,
        audio_path=audio_path,
        audio_dur=5.,
        image_transform=image_transform,
        audio_transform=audio_transform,
        num_classes=num_classes,
        class_labels=class_labels,
        num_mixtures=1
    )


def inverse_normalize(tensor):
    inverse_mean = [-0.485/0.229, -0.456/0.224, -0.406/0.225]
    inverse_std = [1.0/0.229, 1.0/0.224, 1.0/0.225]
    tensor = transforms.Normalize(inverse_mean, inverse_std)(tensor)
    return tensor

def convert_normalize(tensor, new_mean, new_std):
    raw_mean = IMAGENET_DEFAULT_MEAN
    raw_std = IMAGENET_DEFAULT_STD
    # inverse_normalize with raw mean & raw std
    inverse_mean = [-mean/std for mean, std in zip(raw_mean, raw_std)]
    inverse_std = [1.0/std for std in raw_std]
    tensor = transforms.Normalize(inverse_mean, inverse_std)(tensor)
    # normalize with new mean & new std
    tensor = transforms.Normalize(new_mean, new_std)(tensor)
    return tensor
datasets.py
stoneMo/CIGN
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from einops import rearrange
from timm.models.layers import DropPath, to_2tuple, trunc_normal_


class Mlp(nn.Module):

    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class MixerMlp(Mlp):

    def forward(self, x):
        return super().forward(x.transpose(1, 2)).transpose(1, 2)


def hard_softmax(logits, dim):
    y_soft = logits.softmax(dim)
    # Straight through.
    index = y_soft.max(dim, keepdim=True)[1]
    y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
    ret = y_hard - y_soft.detach() + y_soft

    return ret


def gumbel_softmax(logits: torch.Tensor, tau: float = 1, hard: bool = False, dim: int = -1) -> torch.Tensor:
    # _gumbels = (-torch.empty_like(
    #     logits,
    #     memory_format=torch.legacy_contiguous_format).exponential_().log()
    #             )  # ~Gumbel(0,1)
    # more stable https://github.com/pytorch/pytorch/issues/41663
    gumbel_dist = torch.distributions.gumbel.Gumbel(
        torch.tensor(0., device=logits.device, dtype=logits.dtype),
        torch.tensor(1., device=logits.device, dtype=logits.dtype))
    gumbels = gumbel_dist.sample(logits.shape)

    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)

    if hard:
        # Straight through.
        index = y_soft.max(dim, keepdim=True)[1]
        y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
        ret = y_hard - y_soft.detach() + y_soft
    else:
        # Reparametrization trick.
        ret = y_soft
    return ret


class AssignAttention(nn.Module):

    def __init__(self,
                 dim,
                 num_heads=1,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.,
                 proj_drop=0.,
                 hard=True,
                 gumbel=False,
                 gumbel_tau=1.,
                 sum_assign=False,
                 assign_eps=1.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.hard = hard
        self.gumbel = gumbel
        self.gumbel_tau = gumbel_tau
        self.sum_assign = sum_assign
        self.assign_eps = assign_eps

    def get_attn(self, attn, gumbel=None, hard=None):

        if gumbel is None:
            gumbel = self.gumbel

        if hard is None:
            hard = self.hard

        attn_dim = -2
        if gumbel and self.training:
            attn = gumbel_softmax(attn, dim=attn_dim, hard=hard, tau=self.gumbel_tau)
        else:
            if hard:
                attn = hard_softmax(attn, dim=attn_dim)
            else:
                attn = F.softmax(attn, dim=attn_dim)

        return attn

    def forward(self, query, key=None, *, value=None, return_attn=False):
        B, N, C = query.shape
        if key is None:
            key = query
        if value is None:
            value = key
        S = key.size(1)
        # [B, nh, N, C//nh]
        q = rearrange(self.q_proj(query), 'b n (h c)-> b h n c', h=self.num_heads, b=B, n=N, c=C // self.num_heads)
        # [B, nh, S, C//nh]
        k = rearrange(self.k_proj(key), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads)
        # [B, nh, S, C//nh]
        v = rearrange(self.v_proj(value), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads)

        # [B, nh, N, S]
        raw_attn = (q @ k.transpose(-2, -1)) * self.scale

        attn = self.get_attn(raw_attn)
        if return_attn:
            hard_attn = attn.clone()
            soft_attn = self.get_attn(raw_attn, gumbel=False, hard=False)
            attn_dict = {'hard': hard_attn, 'soft': soft_attn}
        else:
            attn_dict = None

        if not self.sum_assign:
            attn = attn / (attn.sum(dim=-1, keepdim=True) + self.assign_eps)
        attn = self.attn_drop(attn)
        assert attn.shape == (B, self.num_heads, N, S)

        # [B, nh, N, C//nh] <- [B, nh, N, S] @ [B, nh, S, C//nh]
        out = rearrange(attn @ v, 'b h n c -> b n (h c)', h=self.num_heads, b=B, n=N, c=C // self.num_heads)

        out = self.proj(out)
        out = self.proj_drop(out)
        return out, attn_dict

    def extra_repr(self):
        return f'num_heads: {self.num_heads}, \n' \
               f'hard: {self.hard}, \n' \
               f'gumbel: {self.gumbel}, \n' \
               f'sum_assign={self.sum_assign}, \n' \
               f'gumbel_tau: {self.gumbel_tau}, \n' \
               f'assign_eps: {self.assign_eps}'


class GroupingBlock(nn.Module):
    """Grouping Block to group similar segments together.

    Args:
        dim (int): Dimension of the input.
        out_dim (int): Dimension of the output.
        num_heads (int): Number of heads in the grouping attention.
        num_output_group (int): Number of output groups.
        norm_layer (nn.Module): Normalization layer to use.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        hard (bool): Whether to use hard or soft assignment. Default: True
        gumbel (bool): Whether to use gumbel softmax. Default: True
        sum_assign (bool): Whether to sum assignment or average. Default: False
        assign_eps (float): Epsilon to avoid divide by zero. Default: 1
        gum_tau (float): Temperature for gumbel softmax. Default: 1
    """

    def __init__(self,
                 *,
                 dim,
                 out_dim,
                 num_heads,
                 num_group_token,
                 num_output_group,
                 norm_layer,
                 mlp_ratio=(0.5, 4.0),
                 hard=True,
                 gumbel=True,
                 sum_assign=False,
                 assign_eps=1.,
                 gumbel_tau=1.):
        super(GroupingBlock, self).__init__()
        self.dim = dim
        self.hard = hard
        self.gumbel = gumbel
        self.sum_assign = sum_assign
        self.num_output_group = num_output_group
        # norm on group_tokens
        self.norm_tokens = norm_layer(dim)
        tokens_dim, channels_dim = [int(x * dim) for x in to_2tuple(mlp_ratio)]
        self.mlp_inter = Mlp(num_group_token, tokens_dim, num_output_group)
        self.norm_post_tokens = norm_layer(dim)
        # norm on x
        self.norm_x = norm_layer(dim)
        self.pre_assign_attn = CrossAttnBlock(
            dim=dim, num_heads=num_heads, mlp_ratio=4, qkv_bias=True, norm_layer=norm_layer, post_norm=True)

        self.assign = AssignAttention(
            dim=dim,
            num_heads=1,
            qkv_bias=True,
            hard=hard,
            gumbel=gumbel,
            gumbel_tau=gumbel_tau,
            sum_assign=sum_assign,
            assign_eps=assign_eps)
        self.norm_new_x = norm_layer(dim)
        self.mlp_channels = Mlp(dim, channels_dim, out_dim)
        if out_dim is not None and dim != out_dim:
            self.reduction = nn.Sequential(norm_layer(dim), nn.Linear(dim, out_dim, bias=False))
        else:
            self.reduction = nn.Identity()

    def extra_repr(self):
        return f'hard={self.hard}, \n' \
               f'gumbel={self.gumbel}, \n' \
               f'sum_assign={self.sum_assign}, \n' \
               f'num_output_group={self.num_output_group}, \n '

    def project_group_token(self, group_tokens):
        """
        Args:
            group_tokens (torch.Tensor): group tokens, [B, S_1, C]

        inter_weight (torch.Tensor): [B, S_2, S_1], S_2 is the new number of
            group tokens, it's already softmaxed along dim=-1

        Returns:
            projected_group_tokens (torch.Tensor): [B, S_2, C]
        """
        # [B, S_2, C] <- [B, S_1, C]
        projected_group_tokens = self.mlp_inter(group_tokens.transpose(1, 2)).transpose(1, 2)
        projected_group_tokens = self.norm_post_tokens(projected_group_tokens)
        return projected_group_tokens

    def forward(self, x, group_tokens, return_attn=False):
        """
        Args:
            x (torch.Tensor): image tokens, [B, L, C]
            group_tokens (torch.Tensor): group tokens, [B, S_1, C]
            return_attn (bool): whether to return attention map

        Returns:
            new_x (torch.Tensor): [B, S_2, C], S_2 is the new number of
                group tokens
        """
        group_tokens = self.norm_tokens(group_tokens)
        x = self.norm_x(x)
        # [B, S_2, C]
        projected_group_tokens = self.project_group_token(group_tokens)
        projected_group_tokens = self.pre_assign_attn(projected_group_tokens, x)
        new_x, attn_dict = self.assign(projected_group_tokens, x, return_attn=return_attn)
        new_x += projected_group_tokens

        new_x = self.reduction(new_x) + self.mlp_channels(self.norm_new_x(new_x))

        return new_x, attn_dict


class GroupingBlock_Han(nn.Module):
    """Grouping Block to group similar segments together.

    Args:
        dim (int): Dimension of the input.
        out_dim (int): Dimension of the output.
        num_heads (int): Number of heads in the grouping attention.
        num_output_group (int): Number of output groups.
        norm_layer (nn.Module): Normalization layer to use.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        hard (bool): Whether to use hard or soft assignment. Default: True
        gumbel (bool): Whether to use gumbel softmax. Default: True
        sum_assign (bool): Whether to sum assignment or average. Default: False
        assign_eps (float): Epsilon to avoid divide by zero. Default: 1
        gum_tau (float): Temperature for gumbel softmax. Default: 1
    """

    def __init__(self,
                 *,
                 dim,
                 out_dim,
                 num_heads,
                 num_group_token,
                 num_output_group,
                 norm_layer,
                 mlp_ratio=(0.5, 4.0),
                 hard=True,
                 gumbel=True,
                 sum_assign=False,
                 assign_eps=1.,
                 gumbel_tau=1.):
        super(GroupingBlock_Han, self).__init__()
        self.dim = dim
        self.hard = hard
        self.gumbel = gumbel
        self.sum_assign = sum_assign
        self.num_output_group = num_output_group
        # norm on group_tokens
        self.norm_tokens = norm_layer(dim)
        tokens_dim, channels_dim = [int(x * dim) for x in to_2tuple(mlp_ratio)]
        self.mlp_inter = Mlp(num_group_token, tokens_dim, num_output_group)
        self.norm_post_tokens = norm_layer(dim)
        # norm on x
        self.norm_x = norm_layer(dim)
        self.pre_assign_attn = CrossAttnBlock(
            dim=dim, num_heads=num_heads, mlp_ratio=4, qkv_bias=True, norm_layer=norm_layer, post_norm=True)

        self.assign = AssignAttention(
            dim=dim,
            num_heads=1,
            qkv_bias=True,
            hard=hard,
            gumbel=gumbel,
            gumbel_tau=gumbel_tau,
            sum_assign=sum_assign,
            assign_eps=assign_eps)
        self.norm_new_x = norm_layer(dim)
        self.mlp_channels = Mlp(dim, channels_dim, out_dim)
        if out_dim is not None and dim != out_dim:
            self.reduction = nn.Sequential(norm_layer(dim), nn.Linear(dim, out_dim, bias=False))
        else:
            self.reduction = nn.Identity()

    def extra_repr(self):
        return f'hard={self.hard}, \n' \
               f'gumbel={self.gumbel}, \n' \
               f'sum_assign={self.sum_assign}, \n' \
               f'num_output_group={self.num_output_group}, \n '

    def project_group_token(self, group_tokens):
        """
        Args:
            group_tokens (torch.Tensor): group tokens, [B, S_1, C]

        inter_weight (torch.Tensor): [B, S_2, S_1], S_2 is the new number of
            group tokens, it's already softmaxed along dim=-1

        Returns:
            projected_group_tokens (torch.Tensor): [B, S_2, C]
        """
        # [B, S_2, C] <- [B, S_1, C]
        projected_group_tokens = self.mlp_inter(group_tokens.transpose(1, 2)).transpose(1, 2)
        projected_group_tokens = self.norm_post_tokens(projected_group_tokens)
        return projected_group_tokens

    def forward(self, x, group_tokens, return_attn=False):
        """
        Args:
            x (torch.Tensor): image tokens, [B, L, C]
            group_tokens (torch.Tensor): group tokens, [B, S_1, C]
            return_attn (bool): whether to return attention map

        Returns:
            new_x (torch.Tensor): [B, S_2, C], S_2 is the new number of
                group tokens
        """
        group_tokens = self.norm_tokens(group_tokens)
        x = self.norm_x(x)
        # [B, S_2, C]
        projected_group_tokens = self.project_group_token(group_tokens)
        projected_group_tokens = self.pre_assign_attn(projected_group_tokens, x)
        new_x, attn_dict = self.assign(projected_group_tokens, x, return_attn=return_attn)
        new_x += projected_group_tokens

        new_x = self.reduction(new_x) + self.mlp_channels(self.norm_new_x(new_x))

        return new_x, attn_dict


class Attention(nn.Module):

    def __init__(self,
                 dim,
                 num_heads,
                 out_dim=None,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.,
                 proj_drop=0.,
                 qkv_fuse=False):
        super().__init__()
        if out_dim is None:
            out_dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5
        self.qkv_fuse = qkv_fuse

        if qkv_fuse:
            self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        else:
            self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
            self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
            self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, out_dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def extra_repr(self):
        return f'num_heads={self.num_heads}, \n' \
               f'qkv_bias={self.scale}, \n' \
               f'qkv_fuse={self.qkv_fuse}'

    def forward(self, query, key=None, *, value=None, mask=None):
        if self.qkv_fuse:
            assert key is None
            assert value is None
            x = query
            B, N, C = x.shape
            S = N
            # [3, B, nh, N, C//nh]
            qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
            # [B, nh, N, C//nh]
            q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
        else:
            B, N, C = query.shape
            if key is None:
                key = query
            if value is None:
                value = key
            S = key.size(1)
            # [B, nh, N, C//nh]
            q = rearrange(self.q_proj(query), 'b n (h c)-> b h n c', h=self.num_heads, b=B, n=N, c=C // self.num_heads)
            # [B, nh, S, C//nh]
            k = rearrange(self.k_proj(key), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads)
            # [B, nh, S, C//nh]
            v = rearrange(self.v_proj(value), 'b n (h c)-> b h n c', h=self.num_heads, b=B, c=C // self.num_heads)

        # [B, nh, N, S]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        if mask is not None:
            attn = attn + mask.unsqueeze(dim=1)
            attn = attn.softmax(dim=-1)
        else:
            attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        assert attn.shape == (B, self.num_heads, N, S)

        # [B, nh, N, C//nh] -> [B, N, C]
        # out = (attn @ v).transpose(1, 2).reshape(B, N, C)
        out = rearrange(attn @ v, 'b h n c -> b n (h c)', h=self.num_heads, b=B, n=N, c=C // self.num_heads)
        out = self.proj(out)
        out = self.proj_drop(out)
        return out


class CrossAttnBlock(nn.Module):

    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm,
                 post_norm=False):
        super().__init__()
        if post_norm:
            self.norm_post = norm_layer(dim)
            self.norm_q = nn.Identity()
            self.norm_k = nn.Identity()
        else:
            self.norm_q = norm_layer(dim)
            self.norm_k = norm_layer(dim)
            self.norm_post = nn.Identity()
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, query, key, *, mask=None):
        x = query
        x = x + self.drop_path(self.attn(self.norm_q(query), self.norm_k(key), mask=mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        x = self.norm_post(x)
        return x


class AttnBlock(nn.Module):

    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
            qkv_fuse=True)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, mask=None):
        x = x + self.drop_path(self.attn(self.norm1(x), mask=mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class GroupingLayer(nn.Module):
    """A Transformer layer with Grouping Block for one stage.

    Args:
        dim (int): Number of input channels.
        num_input_token (int): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer.
            In GroupViT setting, Grouping Block serves as the downsampling layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        group_projector (nn.Module | None, optional): Projector for the grouping layer. Default: None.
        zero_init_group_token (bool): Whether to initialize the grouping token to 0. Default: False.
    """

    def __init__(self,
                 dim,
                 num_input_token,
                 depth,
                 num_heads,
                 num_group_token,
                 mlp_ratio=4.,
                 qkv_bias=True,
                 qk_scale=None,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 norm_layer=nn.LayerNorm,
                 downsample=None,
                 use_checkpoint=False,
                 group_projector=None,
                 zero_init_group_token=False):

        super().__init__()
        self.dim = dim
        self.input_length = num_input_token
        self.depth = depth
        self.use_checkpoint = use_checkpoint
        self.num_group_token = num_group_token
        if num_group_token > 0:
            self.group_token = nn.Parameter(torch.zeros(1, num_group_token, dim))
            if not zero_init_group_token:
                trunc_normal_(self.group_token, std=.02)
        else:
            self.group_token = None

        # build blocks
        self.depth = depth
        blocks = []
        for i in range(depth):
            blocks.append(
                AttnBlock(
                    dim=dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path[i],
                    norm_layer=norm_layer))
        self.blocks = nn.ModuleList(blocks)

        self.downsample = downsample
        self.input_resolution = num_input_token
        self.use_checkpoint = use_checkpoint

        self.group_projector = group_projector

    @property
    def with_group_token(self):
        return self.group_token is not None

    def extra_repr(self):
        return f'dim={self.dim}, \n' \
               f'input_resolution={self.input_resolution}, \n' \
               f'depth={self.depth}, \n' \
               f'num_group_token={self.num_group_token}, \n'

    def split_x(self, x):
        if self.with_group_token:
            return x[:, :-self.num_group_token], x[:, -self.num_group_token:]
        else:
            return x, None

    def concat_x(self, x, group_token=None):
        if group_token is None:
            return x
        return torch.cat([x, group_token], dim=1)

    def forward(self, x, prev_group_token=None, return_attn=False):
        """
        Args:
            x (torch.Tensor): image tokens, [B, L, C]
            prev_group_token (torch.Tensor): group tokens, [B, S_1, C]
            return_attn (bool): whether to return attention maps
        """
        if self.with_group_token:
            group_token = self.group_token.expand(x.size(0), -1, -1)
            if self.group_projector is not None:
                group_token = group_token + self.group_projector(prev_group_token)
        else:
            group_token = None

        B, L, C = x.shape
        cat_x = self.concat_x(x, group_token)
        for blk_idx, blk in enumerate(self.blocks):
            if self.use_checkpoint:
                cat_x = checkpoint.checkpoint(blk, cat_x)
            else:
                cat_x = blk(cat_x)

        x, group_token = self.split_x(cat_x)

        attn_dict = None
        if self.downsample is not None:
            x, attn_dict = self.downsample(x, group_token, return_attn=return_attn)

        return x, group_token, attn_dict


class ModalityTrans(nn.Module):

    def __init__(self,
                dim,
                depth=3,
                num_heads=6,
                mlp_ratio=4.,
                qkv_bias=True,
                qk_scale=None,
                drop=0.,
                attn_drop=0.,
                drop_path=0.,
                norm_layer=nn.LayerNorm,
                out_dim_grouping=512,
                num_heads_grouping=6,
                num_group_tokens=10+25,
                num_output_groups=25,
                hard_assignment=True,
                use_han=True,
                use_grouping=True
                ):

        super(ModalityTrans, self).__init__()

        # build blocks
        self.depth = depth
        blocks = []
        for i in range(depth):
            blocks.append(
                AttnBlock(
                    dim=dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path,
                    norm_layer=norm_layer)
                )

        self.blocks = nn.ModuleList(blocks)

        self.num_group_token = num_group_tokens

        # han encoder
        if use_han:
            # self.han_encoder = Encoder(HANLayer(d_model=dim, nhead=8, dim_feedforward=dim), num_layers=1)
            self.han_encoder = GroupingBlock_Han(
                    dim=dim,
                    out_dim=out_dim_grouping,
                    num_heads=8,
                    num_group_token=10,
                    num_output_group=10,
                    norm_layer=norm_layer,
                    hard=not hard_assignment,
                    gumbel=not hard_assignment
                )
        else:
            self.han_encoder = None

        # grouping block
        if use_grouping:
            self.grouping = GroupingBlock(
                            dim=dim,
                            out_dim=out_dim_grouping,
                            num_heads=num_heads_grouping,
                            num_group_token=num_group_tokens,
                            num_output_group=num_output_groups,
                            norm_layer=norm_layer,
                            hard=hard_assignment,
                            gumbel=hard_assignment
                        )
        else:
            self.grouping = None

    def split_x(self, x):
        return x[:, :-self.num_group_token], x[:, -self.num_group_token:]

    def concat_x(self, x, group_token=None):
        if group_token is None:
            return x
        return torch.cat([x, group_token], dim=1)

    def forward(self, x, group_token=None, x_other=None, return_attn=False):
        """
        Args:
            x (torch.Tensor): modality tokens, [B, L, C]
            x_other (torch.Tensor): another modality tokens, [B, L, C]
            cls_token (torch.Tensor): cls tokens, [B, S_1, C]
            return_attn (bool): whether to return attention maps
        """
        group_token = group_token.expand(x.size(0), -1, -1)

        B, L, C = x.shape
        cat_x = self.concat_x(x, group_token)
        for blk_idx, blk in enumerate(self.blocks):
            cat_x = blk(cat_x)

        x_attn, group_token = self.split_x(cat_x)

        # NOTE: han encoder
        if self.han_encoder is not None and x_other is not None:
            x_attn, _ = self.han_encoder(x_attn, x_other)

        attn_dict = None
        if self.grouping is not None:
            x, attn_dict = self.grouping(x_attn, group_token, return_attn=return_attn)

        return x, attn_dict, x_attn

class UniModalityTrans(nn.Module):

    def __init__(self,
                dim,
                depth=3,
                num_heads=6,
                mlp_ratio=4.,
                qkv_bias=True,
                qk_scale=None,
                drop=0.,
                attn_drop=0.,
                drop_path=0.,
                norm_layer=nn.LayerNorm,
                out_dim_grouping=512,
                num_heads_grouping=6,
                num_group_tokens=10+25,
                num_output_groups=25,
                hard_assignment=True,
                use_han=True
                ):

        super(UniModalityTrans, self).__init__()

        # build blocks
        self.depth = depth
        blocks = []
        for i in range(depth):
            blocks.append(
                AttnBlock(
                    dim=dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path,
                    norm_layer=norm_layer)
                )

        self.blocks = nn.ModuleList(blocks)

        self.num_group_token = num_group_tokens

        # han encoder
        if use_han:
            # self.han_encoder = Encoder(HANLayer(d_model=dim, nhead=8, dim_feedforward=dim), num_layers=1)
            self.han_encoder = GroupingBlock_Han(
                    dim=dim,
                    out_dim=out_dim_grouping,
                    num_heads=8,
                    num_group_token=10,
                    num_output_group=10,
                    norm_layer=norm_layer,
                    hard=not hard_assignment,
                    gumbel=not hard_assignment
                )
        else:
            self.han_encoder = None

        # grouping block
        self.grouping = GroupingBlock(
                    dim=dim,
                    out_dim=out_dim_grouping,
                    num_heads=num_heads_grouping,
                    num_group_token=num_group_tokens,
                    num_output_group=num_output_groups,
                    norm_layer=norm_layer,
                    hard=hard_assignment,
                    gumbel=hard_assignment
                )

    def split_x(self, x):
        return x[:, :-self.num_group_token], x[:, -self.num_group_token:]

    def concat_x(self, x, group_token=None):
        if group_token is None:
            return x
        return torch.cat([x, group_token], dim=1)

    def forward(self, x, group_token=None, x_other=None, return_attn=False):
        """
        Args:
            x (torch.Tensor): modality tokens, [B, L, C]
            x_other (torch.Tensor): another modality tokens, [B, L, C]
            cls_token (torch.Tensor): cls tokens, [B, S_1, C]
            return_attn (bool): whether to return attention maps
        """
        group_token = group_token.expand(x.size(0), -1, -1)

        B, L, C = x.shape
        cat_x = self.concat_x(x, group_token)

        for blk_idx, blk in enumerate(self.blocks):
            cat_x = blk(cat_x)

        x_attn, group_token = self.split_x(cat_x)

        # NOTE: han encoder
        if self.han_encoder is not None and x_other is not None:
            x_attn, _ = self.han_encoder(x_attn, x_other)

        attn_dict = None
        if self.grouping is not None:
            x, attn_dict = self.grouping(x_attn, group_token, return_attn=return_attn)

        return x, attn_dict, x_attn
grouping.py
stoneMo/CIGN
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.models import resnet18
import torch.distributed as dist
from timm.models.layers import DropPath, to_2tuple, trunc_normal_

from dino.hubconf import *
from grouping import ModalityTrans


class CIGN(nn.Module):
    def __init__(self, tau, dim, dropout_img, dropout_aud, args):
        super(CIGN, self).__init__()
        self.tau = tau

        # Vision model
        self.imgnet = resnet18(pretrained=True)
        self.imgnet.avgpool = nn.Identity()
        self.imgnet.fc = nn.Identity()
        self.img_proj = nn.Conv2d(512, dim, kernel_size=(1, 1))
        self.img_dropout = nn.Dropout(p=dropout_img)

        # Audio model
        self.audnet = resnet18()
        self.audnet.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        self.audnet.avgpool = nn.AdaptiveMaxPool2d((1, 1))
        # self.audnet.avgpool = nn.Identity()
        self.audnet.fc = nn.Identity()
        self.aud_proj = nn.Linear(512, dim)
        # self.aud_proj = nn.Conv2d(512, dim, kernel_size=(1, 1))
        self.aud_dropout = nn.Dropout(p=dropout_aud)

        # Initialize weights (except pretrained visual model)
        for net in [self.audnet, self.img_proj, self.aud_proj]:
            for m in net.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.kaiming_normal_(
                        m.weight, mode='fan_out', nonlinearity='relu')
                elif isinstance(m, nn.Linear):
                    nn.init.trunc_normal_(
                        m.weight, mean=0.0, std=0.01)
                elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                    nn.init.normal_(m.weight, mean=1, std=0.02)
                    nn.init.constant_(m.bias, 0)

        # hard or soft assignment
        self.unimodal_assgin = args.attn_assign
        unimodal_hard_assignment = True if args.attn_assign == 'hard' else False

        # learnable tokens
        self.num_class = args.num_class
        self.train_stage = args.train_stage
        if self.train_stage == 0:
            self.av_token = nn.Parameter(torch.zeros(args.num_class//1, args.dim))
        elif self.train_stage == 1:
            self.av_token = nn.Parameter(torch.zeros(args.num_class//2, args.dim), requires_grad=False)
            self.av_token_1 = nn.Parameter(torch.zeros(args.num_class//2, args.dim))
        elif self.train_stage == 2:
            self.av_token = nn.Parameter(torch.zeros(args.num_class//3, args.dim), requires_grad=False)
            self.av_token_1 = nn.Parameter(torch.zeros(args.num_class//3, args.dim), requires_grad=False)
            self.av_token_2 = nn.Parameter(torch.zeros(args.num_class//3, args.dim))
        elif self.train_stage == 3:
            self.av_token = nn.Parameter(torch.zeros(args.num_class//4, args.dim))
            self.av_token_1 = nn.Parameter(torch.zeros(args.num_class//4, args.dim))
            self.av_token_2 = nn.Parameter(torch.zeros(args.num_class//4, args.dim))
            self.av_token_3 = nn.Parameter(torch.zeros(args.num_class//4, args.dim))
        
        self.test_stage = args.test_stage

        # uni-modal encoder
        self.audio_encoder = ModalityTrans(
                            args.dim,
                            depth=args.depth_aud,
                            num_heads=8,
                            mlp_ratio=4.,
                            qkv_bias=True,
                            qk_scale=None,
                            drop=0.,
                            attn_drop=0.,
                            drop_path=0.1,
                            norm_layer=nn.LayerNorm,
                            out_dim_grouping=args.dim,
                            num_heads_grouping=8,
                            num_group_tokens=args.num_class,
                            num_output_groups=args.num_class,
                            hard_assignment=unimodal_hard_assignment,
                            use_han=True
                        )

        self.visual_encoder = ModalityTrans(
                            args.dim,
                            depth=args.depth_vis,
                            num_heads=8,
                            mlp_ratio=4.,
                            qkv_bias=True,
                            qk_scale=None,
                            drop=0.,
                            attn_drop=0.,
                            drop_path=0.1,
                            norm_layer=nn.LayerNorm,
                            out_dim_grouping=args.dim,
                            num_heads_grouping=8,
                            num_group_tokens=args.num_class,
                            num_output_groups=args.num_class,
                            hard_assignment=unimodal_hard_assignment,
                            use_han=False
                        )

        # prediction heads
        self.fc_prob_a = nn.Linear(args.dim, 1)
        self.fc_prob_v = nn.Linear(args.dim, 1)
        self.fc_cls = nn.Linear(args.dim, args.num_class)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def max_xmil_loss(self, img, aud):
        B = img.shape[0]
        Slogits = torch.einsum('nchw,mc->nmhw', img, aud) / self.tau
        logits = Slogits.flatten(-2, -1).max(dim=-1)[0]
        labels = torch.arange(B).long().to(img.device)
        loss = F.cross_entropy(logits, labels) + F.cross_entropy(logits.permute(1, 0), labels)
        return loss, Slogits

    def cls_token_loss(self, cls_prob):
        cls_target = torch.arange(0,self.num_class).long().to(cls_prob.device)
        loss = F.cross_entropy(cls_prob, cls_target)
        return loss, cls_target

    def cls_pred_loss(self, prob, target):
        loss = F.binary_cross_entropy(prob, target)
        return loss

    def forward(self, image, audio, mode='train', cls_target=None):
           
        cls_target_v = cls_target.copy()
        cls_target_a = cls_target.copy()

        # Image
        img = self.imgnet(image).unflatten(1, (512, 7, 7))
        img = self.img_dropout(img)
        img = self.img_proj(img)
        img = nn.functional.normalize(img, dim=1)

        # Audio
        aud = self.audnet(audio)
        aud = self.aud_dropout(aud)
        aud = self.aud_proj(aud)
        aud = nn.functional.normalize(aud, dim=1)
         
        if self.test_stage == 0:
            av_token = self.av_token
        elif self.test_stage == 1:
            av_token = torch.cat([self.av_token, self.av_token_1], dim=0)
        elif self.test_stage == 2:
            av_token = torch.cat([self.av_token, self.av_token_1, self.av_token_2], dim=0)
        elif self.test_stage == 3:
            av_token = torch.cat([self.av_token, self.av_token_1, self.av_token_2, self.av_token_3], dim=0)

        xv = img.flatten(2,3).permute(0,2,1)
        xv, _, _ = self.visual_encoder(xv, av_token, return_attn=True)
        xa = aud.unsqueeze(1)
        xa, _, _ = self.audio_encoder(xa, av_token, return_attn=True)

        av_cls_prob = self.fc_cls(av_token)                                      
        a_prob = torch.sigmoid(self.fc_prob_a(xa))                                  
        a_pred_prob = a_prob.sum(dim=-1)                                            
        v_prob = torch.sigmoid(self.fc_prob_v(xv))                                  
        v_pred_prob = v_prob.sum(dim=-1)                                            

        cls_token_loss, av_cls_target = self.cls_token_loss(av_cls_prob)
        cls_pred_loss = self.cls_pred_loss(v_pred_prob, cls_target_v) + self.cls_pred_loss(a_pred_prob, cls_target_a)
        
        return a_pred_prob, v_pred_prob, cls_token_loss, cls_pred_loss, av_cls_prob, av_cls_target
        
model.py
stoneMo/CIGN
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import utils
import numpy as np
import argparse
from model import EZVSL, MoVSL, VSGN
from datasets import get_test_dataset, inverse_normalize
import cv2

import torch.multiprocessing as mp
import torch.distributed as dist

from object_model import DINO


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str, default='./checkpoints', help='path to save trained model weights')
    parser.add_argument('--experiment_name', type=str, default='movsl_vggss', help='experiment name (experiment folder set to "args.model_dir/args.experiment_name)"')
    parser.add_argument('--save_visualizations', action='store_true', help='Set to store all VSL visualizations (saved in viz directory within experiment folder)')

    # Dataset
    parser.add_argument('--testset', default='flickr', type=str, help='testset (flickr or vggss)')
    parser.add_argument('--test_data_path', default='', type=str, help='Root directory path of data')
    parser.add_argument('--test_gt_path', default='', type=str)
    parser.add_argument('--batch_size', default=1, type=int, help='Batch Size')
    parser.add_argument('--num_class', default=37, type=int)

    # mo-vsl hyper-params
    parser.add_argument('--model', default='movsl')
    parser.add_argument('--out_dim', default=512, type=int)
    parser.add_argument('--num_negs', default=None, type=int)
    parser.add_argument('--tau', default=0.03, type=float, help='tau')

    parser.add_argument('--attn_assign', type=str, default='soft', help="type of audio grouping assignment")
    parser.add_argument('--dim', type=int, default=512, help='dimensionality of features')
    parser.add_argument('--depth_aud', type=int, default=3, help='depth of audio transformers')
    parser.add_argument('--depth_vis', type=int, default=3, help='depth of visual transformers')

    # evaluation parameters
    parser.add_argument("--dropout_img", type=float, default=0, help="dropout for image")
    parser.add_argument("--dropout_aud", type=float, default=0, help="dropout for audio")

    parser.add_argument('--m_img', default=1.0, type=float, metavar='M', help='momentum for imgnet')
    parser.add_argument('--m_aud', default=1.0, type=float, metavar='M', help='momentum for audnet')

    parser.add_argument('--use_momentum', action='store_true')
    parser.add_argument('--relative_prediction', action='store_true')
    parser.add_argument('--use_mom_eval', action='store_true')
    parser.add_argument('--pred_size', default=0.5, type=float)
    parser.add_argument('--pred_thr', default=0.5, type=float)

    parser.add_argument('--object_model_type', default='res18', type=str)

    # Distributed params
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--gpu', type=int, default=None)
    parser.add_argument('--world_size', type=int, default=1)
    parser.add_argument('--rank', type=int, default=0)
    parser.add_argument('--node', type=str, default='localhost')
    parser.add_argument('--port', type=int, default=12345)
    parser.add_argument('--dist_url', type=str, default='tcp://localhost:12345')
    parser.add_argument('--multiprocessing_distributed', action='store_true')

    return parser.parse_args()


def main(args):
    mp.set_start_method('spawn')
    args.dist_url = f'tcp://{args.node}:{args.port}'
    print('Using url {}'.format(args.dist_url))

    ngpus_per_node = torch.cuda.device_count()
    if args.multiprocessing_distributed:
        args.world_size = ngpus_per_node
        mp.spawn(main_worker,
                 nprocs=ngpus_per_node,
                 args=(ngpus_per_node, args))

    else:
        main_worker(args.gpu, ngpus_per_node, args)


def main_worker(local_rank, ngpus_per_node, args):
    args.gpu = local_rank
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # Model dir
    model_dir = os.path.join(args.model_dir, args.experiment_name)
    viz_dir = os.path.join(model_dir, 'viz')
    os.makedirs(viz_dir, exist_ok=True)

    # Setup distributed environment
    if args.multiprocessing_distributed:
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + local_rank
            print(args.dist_url, args.world_size, args.rank)
        dist.init_process_group(backend='nccl', init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)
        torch.distributed.barrier()

    # Create model
    if args.model.lower() == 'ezvsl':
        audio_visual_model = EZVSL(args.tau, args.out_dim, args.dropout_img, args.dropout_aud)
    elif args.model.lower() == 'vsgn':
        audio_visual_model = VSGN(args.tau, args.out_dim, args.dropout_img, args.dropout_aud, args)
    elif args.model.lower() == 'movsl':
        audio_visual_model = MoVSL(args.tau, args.out_dim, args.dropout_img, args.dropout_aud, args.m_img, args.m_aud, args.use_mom_eval, num_neg=args.num_negs)
    else:
        raise ValueError

    if args.object_model_type =='dino':
        object_saliency_model = DINO(arch_type='vitb16')
    else:
        from torchvision.models import resnet18
        object_saliency_model = resnet18(pretrained=True)
        object_saliency_model.avgpool = nn.Identity()
        object_saliency_model.fc = nn.Sequential(
            nn.Unflatten(1, (512, 7, 7)),
            NormReducer(dim=1),
            Unsqueeze(1)
        )
        # object_saliency_model.fc = nn.Unflatten(1, (512, 7, 7))

    if not torch.cuda.is_available():
        print('using CPU, this will be slow')
    else:
        if args.gpu is not None:
            torch.cuda.set_device(args.gpu)
            audio_visual_model.cuda(args.gpu)
            object_saliency_model.cuda(args.gpu)
            if args.multiprocessing_distributed:
                audio_visual_model = torch.nn.parallel.DistributedDataParallel(audio_visual_model, device_ids=[args.gpu])
                object_saliency_model = torch.nn.parallel.DistributedDataParallel(object_saliency_model, device_ids=[args.gpu])

    # Load weights
    ckp_fn = os.path.join(model_dir, 'best.pth')
    if os.path.exists(ckp_fn):
        ckp = torch.load(ckp_fn, map_location='cpu')
        audio_visual_model.load_state_dict({k.replace('module.', ''): ckp['model'][k] for k in ckp['model']})
        print(f'loaded from {os.path.join(model_dir, "best.pth")}')
    else:
        print(f"Checkpoint not found: {ckp_fn}")

    # Dataloader
    testdataset = get_test_dataset(args)
    if args.multiprocessing_distributed:
        sampler = torch.utils.data.DistributedSampler(testdataset, num_replicas=ngpus_per_node, rank=args.rank, shuffle=False)
    else:
        sampler = torch.utils.data.SequentialSampler(testdataset)
    testdataloader = DataLoader(testdataset, batch_size=args.batch_size, sampler=sampler, num_workers=args.workers)
    print("Loaded dataloader.")

    validate(testdataloader, audio_visual_model, object_saliency_model, model_dir, args)


@torch.no_grad()
def validate(testdataloader, audio_visual_model, object_saliency_model, model_dir, args):
    audio_visual_model.train(False)
    object_saliency_model.train(False)

    evaluator_av = utils.EvaluatorFull(default_conf_thr=0.5, pred_size=args.pred_size, pred_thr=args.pred_thr, results_dir=f"{model_dir}/av")
    evaluator_obj = utils.EvaluatorFull(default_conf_thr=0.5, pred_size=args.pred_size,  pred_thr=args.pred_thr, results_dir=f"{model_dir}/obj")
    evaluator_av_obj = utils.EvaluatorFull(default_conf_thr=0., pred_size=args.pred_size,  pred_thr=args.pred_thr, results_dir=f"{model_dir}/av_obj")
    for step, (image, spec, bboxes, name) in enumerate(testdataloader):
        if args.gpu is not None:
            spec = spec.cuda(args.gpu, non_blocking=True)
            image = image.cuda(args.gpu, non_blocking=True)

        # Compute S_AVL
        heatmap_av = audio_visual_model(image.float(), spec.float(), mode='test')[1].unsqueeze(1)
        heatmap_av = F.interpolate(heatmap_av, size=(224, 224), mode='bicubic', align_corners=True)
        heatmap_av = heatmap_av.data.cpu().numpy()

        # Compute S_OBJ
        img_feat = object_saliency_model(image)
        heatmap_obj = F.interpolate(img_feat, size=(224, 224), mode='bicubic', align_corners=True)
        heatmap_obj = heatmap_obj.data.cpu().numpy()

        av_min, av_max = -1. / args.tau, 1. / args.tau
        obj_min, obj_max = 0., 2.5
        min_max_norm = lambda x, xmin, xmax: (x - xmin) / (xmax - xmin)

        # Compute eval metrics and save visualizations
        for i in range(spec.shape[0]):
            gt_map = bboxes['gt_map'][i].data.cpu().numpy()
            bb = bboxes['bboxes'][i]
            bb = bb[bb[:, 0] >= 0].numpy().tolist()

            n = heatmap_av[i, 0].size
            scores_av = min_max_norm(heatmap_av[i, 0], av_min, av_max)
            scores_obj = min_max_norm(heatmap_obj[i, 0], obj_min, obj_max)
            scores_av_obj = scores_av * args.alpha + scores_obj * (1 - args.alpha)

            conf_av = np.sort(scores_av.flatten())[-n//4:].mean()
            conf_obj = np.sort(scores_obj.flatten())[-n//4:].mean()
            conf_av_obj = np.sort(scores_av_obj.flatten())[-n//4:].mean()

            if args.relative_prediction:
                pred_av = utils.normalize_img(scores_av)
                pred_obj = utils.normalize_img(scores_obj)
                pred_av_obj = utils.normalize_img(scores_av_obj)

                thr_av = np.sort(pred_av.flatten())[int(n * args.pred_size)]
                thr_obj = np.sort(pred_obj.flatten())[int(n * args.pred_size)]
                thr_av_obj = np.sort(pred_av_obj.flatten())[int(n * args.pred_size)]
            else:
                pred_av = scores_av
                pred_obj = scores_obj
                pred_av_obj = scores_av_obj

                thr_av = thr_obj = thr_av_obj = args.pred_thr

            evaluator_av.update(bb, gt_map, conf_av, pred_av, thr_av, name[i])
            evaluator_obj.update(bb, gt_map, conf_obj, pred_obj, thr_obj, name[i])
            evaluator_av_obj.update(bb, gt_map, conf_av_obj, pred_av_obj, thr_av_obj, name[i])

            if args.save_visualizations:
                evaluator_av.save_viz(image[i], bb, pred_av, name[i])
                evaluator_obj.save_viz(image[i], bb, pred_obj, name[i])
                evaluator_av_obj.save_viz(image[i], bb, pred_av_obj, name[i])

        print(f'{step+1}/{len(testdataloader)}: AV+OGL-Prec@30={evaluator_av_obj.precision_at_30():.3f} AVL-Prec@30={evaluator_av.precision_at_30():.3f} OGL-Prec@30={evaluator_obj.precision_at_30():.3f}')

    evaluator_av.save_results()
    evaluator_obj.save_results()
    evaluator_av_obj.save_results()

    print('='*20 + ' AVL ' + '='*20)
    stats_av = evaluator_av.finalize_stats()
    print('\n'.join([f' - {k}: {stats_av[k]}' for k in sorted(stats_av.keys()) if stats_av[k] is not np.nan]))

    print('='*20 + ' OGL ' + '='*20)
    stats_obj = evaluator_obj.finalize_stats()
    print('\n'.join([f' - {k}: {stats_obj[k]}' for k in sorted(stats_obj.keys()) if stats_obj[k] is not np.nan]))

    print('='*20 + ' AV+OGL ' + '='*20)
    stats_av_obj = evaluator_av_obj.finalize_stats()
    print('\n'.join([f' - {k}: {stats_av_obj[k]}' for k in sorted(stats_av_obj.keys()) if stats_av_obj[k] is not np.nan]))


class NormReducer(nn.Module):
    def __init__(self, dim):
        super(NormReducer, self).__init__()
        self.dim = dim

    def forward(self, x):
        return x.abs().mean(self.dim)


class Unsqueeze(nn.Module):
    def __init__(self, dim):
        super(Unsqueeze, self).__init__()
        self.dim = dim

    def forward(self, x):
        return x.unsqueeze(self.dim)


if __name__ == "__main__":
    main(get_arguments())

test.py
stoneMo/CIGN
from multiprocessing import reduction
import os
import argparse
import builtins
import sys
import time
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import multiprocessing as mp
import torch.distributed as dist
from tensorboardX import SummaryWriter
from sklearn.metrics import precision_score, recall_score, f1_score

import utils
from model import CIGN
from datasets import get_train_dataset, get_test_dataset


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str, default='./checkpoints', help='path to save trained model weights')
    parser.add_argument('--experiment_name', type=str, default='cign_vggss', help='experiment name (used for checkpointing and logging)')
    parser.add_argument('--experiment_name_pre', type=str, default='cign_vggss', help='experiment name of previous stage (used for checkpointing and logging)')

    # Data params
    parser.add_argument('--trainset', default='vggss', type=str, help='trainset (flickr or vggss)')
    parser.add_argument('--testset', default='vggss', type=str, help='testset,(flickr or vggss)')
    parser.add_argument('--train_data_path', default='', type=str, help='Root directory path of train data')
    parser.add_argument('--test_data_path', default='', type=str, help='Root directory path of test data')
    parser.add_argument('--test_gt_path', default='', type=str)
    parser.add_argument('--num_test_samples', default=-1, type=int)
    parser.add_argument('--num_class', default=37, type=int)

    parser.add_argument('--sup_train_data_path', default='', type=str, help='Root directory path of train data w. labels')
    parser.add_argument('--sup_train_gt_path', default='', type=str)
    parser.add_argument('--use_supervised_data', action='store_true')

    # mo-vsl hyper-params
    parser.add_argument('--model', default='cign')
    parser.add_argument('--imgnet_type', default='vitb8')
    parser.add_argument('--audnet_type', default='vitb8')

    parser.add_argument('--out_dim', default=512, type=int)
    parser.add_argument('--num_negs', default=None, type=int)
    parser.add_argument('--tau', default=0.03, type=float, help='tau')

    parser.add_argument('--attn_assign', type=str, default='soft', help="type of audio grouping assignment")
    parser.add_argument('--dim', type=int, default=512, help='dimensionality of features')
    parser.add_argument('--depth_aud', type=int, default=3, help='depth of audio transformers')
    parser.add_argument('--depth_vis', type=int, default=3, help='depth of visual transformers')

    # training/evaluation parameters
    parser.add_argument("--epochs", type=int, default=20, help="number of epochs")
    parser.add_argument('--batch_size', default=128, type=int, help='Batch Size')
    parser.add_argument("--lr_schedule", default='cte', help="learning rate schedule")
    parser.add_argument("--init_lr", type=float, default=0.0001, help="initial learning rate")
    parser.add_argument("--warmup_epochs", type=int, default=0, help="warmup epochs")
    parser.add_argument("--seed", type=int, default=12345, help="random seed")
    parser.add_argument('--weight_decay', type=float, default=0, help='Weight Decay')
    parser.add_argument("--clip_norm", type=float, default=0, help="gradient clip norm")
    parser.add_argument("--dropout_img", type=float, default=0, help="dropout for image")
    parser.add_argument("--dropout_aud", type=float, default=0, help="dropout for audio")
    
    parser.add_argument("--train_stage", type=int, default=0, help="training stage")
    parser.add_argument("--test_stage", type=int, default=0, help="testing stage")
    parser.add_argument('--resume_av_token', action='store_true')

    parser.add_argument('--m_img', default=0.9, type=float, metavar='M', help='momentum for imgnet')
    parser.add_argument('--m_aud', default=0.9, type=float, metavar='M', help='momentum for audnet')
    parser.add_argument('--use_momentum', action='store_true')
    parser.add_argument('--use_mom_eval', action='store_true')

    # Distributed params
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--gpu', type=int, default=None)
    parser.add_argument('--world_size', type=int, default=1)
    parser.add_argument('--rank', type=int, default=0)
    parser.add_argument('--node', type=str, default='localhost')
    parser.add_argument('--port', type=int, default=12345)
    parser.add_argument('--dist_url', type=str, default='tcp://localhost:12345')
    parser.add_argument('--multiprocessing_distributed', action='store_true')

    return parser.parse_args()


def main(args):
    mp.set_start_method('spawn')
    args.dist_url = f'tcp://{args.node}:{args.port}'
    print('Using url {}'.format(args.dist_url))

    ngpus_per_node = torch.cuda.device_count()
    if args.multiprocessing_distributed:
        args.world_size = ngpus_per_node
        mp.spawn(main_worker,
                 nprocs=ngpus_per_node,
                 args=(ngpus_per_node, args))

    else:
        main_worker(args.gpu, ngpus_per_node, args)


def main_worker(gpu, ngpus_per_node, args):
    args.gpu = gpu

    if args.gpu is not None:
        print("Use GPU: {} for training".format(args.gpu))

    # Setup distributed environment
    if args.multiprocessing_distributed:
        if args.dist_url == "env://" and args.rank == -1:
            args.rank = int(os.environ["RANK"])
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + gpu
        dist.init_process_group(backend='nccl', init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)
        torch.distributed.barrier()

    # Create model dir
    model_dir = os.path.join(args.model_dir, args.experiment_name)
    os.makedirs(model_dir, exist_ok=True)
    utils.save_json(vars(args), os.path.join(model_dir, 'configs.json'), sort_keys=True, save_pretty=True)

    # tb writers
    tb_writer = SummaryWriter(model_dir)

    # logger
    log_fn = f"{model_dir}/train.log"
    def print_and_log(*content, **kwargs):
        # suppress printing if not first GPU on each node
        if args.multiprocessing_distributed and (args.gpu != 0 or args.rank != 0):
            return
        msg = ' '.join([str(ct) for ct in content])
        sys.stdout.write(msg+'\n')
        sys.stdout.flush()
        with open(log_fn, 'a') as f:
            f.write(msg+'\n')
    builtins.print = print_and_log

    # Create model
    if args.model.lower() == 'CIGN':
        model = CIGN(args.tau, args.out_dim, args.dropout_img, args.dropout_aud, args)
    else:
        raise ValueError

    if not torch.cuda.is_available():
        print('using CPU, this will be slow')
    elif args.multiprocessing_distributed:
        if args.gpu is not None:
            torch.cuda.set_device(args.gpu)
            model.cuda(args.gpu)
            # When using a single GPU per process and per
            # DistributedDataParallel, we need to divide the batch size
            # ourselves based on the total number of GPUs we have
            args.batch_size = int(args.batch_size / args.world_size)
            args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
    elif args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model.cuda(args.gpu)
    print(model)

    # Optimizer
    optimizer, _ = utils.build_optimizer_and_scheduler_adam(model, args)

    # load av tokens trained from previous stage
    if args.resume_av_token:
        model_dir_pre = os.path.join(args.model_dir, args.experiment_name_pre)
        ckp = torch.load(os.path.join(model_dir_pre, 'best.pth'), map_location='cpu')
        pretrained_tokens = {k: v for k, v in ckp['model'].items() if k.startswith('module.av_token')}
        print('pretrained_tokens:', pretrained_tokens.keys())
        model_dict = model.state_dict()
        model_dict.update(pretrained_tokens)
        model.load_state_dict(model_dict, strict=False)

    # Resume current stage if possible
    start_epoch, best_aacc, best_vacc, best_avacc = 0, 0., 0., 0.
    if os.path.exists(os.path.join(model_dir, 'best.pth')):
        ckp = torch.load(os.path.join(model_dir, 'best.pth'), map_location='cpu')
        start_epoch, best_aacc, best_vacc, best_avacc = ckp['epoch'], ckp['best_AACC'], ckp['best_VACC'], ckp['best_AVACC']
        model.load_state_dict(ckp['model'])
        optimizer.load_state_dict(ckp['optimizer'])
        print(f'loaded from {os.path.join(model_dir, "best.pth")}')

    # Dataloaders
    traindataset = get_train_dataset(args)
    train_sampler = None
    if args.multiprocessing_distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(traindataset)
    train_loader = torch.utils.data.DataLoader(
        traindataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=False, sampler=train_sampler, drop_last=True,
        persistent_workers=args.workers > 0)

    valdataset = get_test_dataset(args, mode='val')
    val_loader = torch.utils.data.DataLoader(
        valdataset, batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=False, drop_last=False,
        persistent_workers=args.workers > 0)

    print("Loaded dataloader.")

    # =============================================================== #
    # Training loop
    for stage_idx in range(args.test_stage+1):
        a_acc, v_acc, av_acc = validate(val_loader, model, args, test_stage=stage_idx)
        print(f'Test Stage (epoch {start_epoch}): {stage_idx}')
        print(f'Audio Acc (epoch {start_epoch}): {a_acc}')
        print(f'Visual Acc (epoch {start_epoch}): {v_acc}')
        print(f'Audio-visual Acc (epoch {start_epoch}): {av_acc}')

    print(f'best_AACC: {best_aacc}')
    print(f'best_VACC: {best_vacc}')
    print(f'best_AVACC: {best_avacc}')

    metric_list = [[] for _ in range(3)]

    for epoch in range(start_epoch, args.epochs):
        if args.multiprocessing_distributed:
            train_loader.sampler.set_epoch(epoch)

        # Train
        train(train_loader, model, optimizer, epoch, args, tb_writer)

        # Evaluate
        a_acc_list, v_acc_list, av_acc_list = [], [], []
        for stage_idx in range(args.test_stage+1):
            a_acc, v_acc, av_acc = validate(val_loader, model, args, test_stage=stage_idx)
            print(f'Test Stage (epoch {epoch+1}): {stage_idx}')
            print(f'Audio Acc (epoch {epoch+1}): {a_acc}')
            print(f'Visual Acc (epoch {epoch+1}): {v_acc}')
            print(f'Audio-visual Acc (epoch {epoch+1}): {av_acc}')

            tb_writer.add_scalar(f'Stage {stage_idx} Audio Acc', a_acc, epoch)
            tb_writer.add_scalar(f'Stage {stage_idx} Visual Acc', v_acc, epoch)
            tb_writer.add_scalar(f'Stage {stage_idx} Audio-visual Acc', av_acc, epoch)

            a_acc_list.append(a_acc)
            v_acc_list.append(v_acc)
            av_acc_list.append(av_acc)

        a_acc, v_acc, av_acc = np.mean(a_acc_list), np.mean(v_acc_list), np.mean(av_acc_list)

        if av_acc >= best_avacc:
           best_aacc, best_vacc, best_avacc = a_acc, v_acc, av_acc
        
        print(f'best_AACC: {best_aacc}')
        print(f'best_VACC: {best_vacc}')
        print(f'best_AVACC: {best_avacc}')

        metric_list[0].append(a_acc)
        metric_list[1].append(v_acc)
        metric_list[2].append(av_acc)

        # Checkpoint
        if args.rank == 0:
            ckp = {'model': model.state_dict(),
                   'optimizer': optimizer.state_dict(),
                   'epoch': epoch+1,
                   'best_AACC': best_aacc,
                   'best_VACC': best_vacc,
                   'best_AVACC': best_avacc}
            torch.save(ckp, os.path.join(model_dir, 'latest.pth'))
            if av_acc == best_avacc:
                torch.save(ckp, os.path.join(model_dir, 'best.pth'))
            print(f"Model saved to {model_dir}")

    np.save(os.path.join(model_dir, 'metrics.npy'), np.array(metric_list))


    # ============Test============

    # Load weights
    ckp_fn = os.path.join(model_dir, 'best.pth')
    ckp = torch.load(ckp_fn, map_location='cpu')
    model.load_state_dict(ckp['model'])
    print(f'loaded from {ckp_fn}')

    testdataset = get_test_dataset(args, mode='test')
    test_loader = torch.utils.data.DataLoader(
        testdataset, batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=False, drop_last=False,
        persistent_workers=args.workers > 0)

    a_acc_list = []
    v_acc_list = []
    av_acc_list = []
    
    for stage_idx in range(args.test_stage+1):
        a_acc, v_acc, av_acc = validate(test_loader, model, args, test_stage=stage_idx)
        print(f'Test Stage (epoch {epoch+1}): {stage_idx}')
        print(f'Audio Acc (epoch {epoch+1}): {a_acc}')
        print(f'Visual Acc (epoch {epoch+1}): {v_acc}')
        print(f'Audio-visual Acc (epoch {epoch+1}): {av_acc}')

        tb_writer.add_scalar(f'Stage {stage_idx} Audio Acc', a_acc, epoch)
        tb_writer.add_scalar(f'Stage {stage_idx} Visual Acc', v_acc, epoch)
        tb_writer.add_scalar(f'Stage {stage_idx} Audio-visual Acc', av_acc, epoch)

        a_acc_list.append(a_acc)
        v_acc_list.append(v_acc)
        av_acc_list.append(av_acc)

    a_acc, v_acc, av_acc = np.mean(a_acc_list), np.mean(v_acc_list), np.mean(av_acc_list)

    print(f'Average_AACC: {a_acc}')
    print(f'Average_VACC: {v_acc}')
    print(f'Average_AVACC: {av_acc}')


def cal_precision_recall_token(target, output):
    precison = precision_score(y_true=target, y_pred=output, average='weighted')
    recall = recall_score(y_true=target, y_pred=output, average='weighted')
    f1 = f1_score(y_true=target, y_pred=output, average='weighted')
    return precison, recall, f1


def train(train_loader, model, optimizer, epoch, args, writer):
    model.train()
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    loss_mtr = AverageMeter('Loss', ':.3f')
    loss_loc_mtr = AverageMeter('Loc Loss', ':.3f')
    loss_token_mtr = AverageMeter('Token Loss', ':.3f')
    loss_pred_mtr = AverageMeter('Pred Loss', ':.3f')

    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, loss_mtr, loss_loc_mtr, loss_token_mtr, loss_pred_mtr],
        prefix="Epoch: [{}]".format(epoch),
    )

    end = time.time()
    for i, (image, spec, anno, _) in enumerate(train_loader):
        data_time.update(time.time() - end)
        global_step = i + len(train_loader) * epoch
        utils.adjust_learning_rate(optimizer, epoch + i / len(train_loader), args)

        if args.gpu is not None:
            spec = spec.cuda(args.gpu, non_blocking=True)
            image = image.cuda(args.gpu, non_blocking=True)
            label = anno['class'].cuda(args.gpu, non_blocking=True)

        _, _, cls_token_loss, cls_pred_loss, av_cls_prob, av_cls_target = model(image.float(), spec.float(), cls_target=label, mode='train')
        loss = cls_token_loss + cls_pred_loss

        loss_mtr.update(loss.item(), image.shape[0])
        loss_token_mtr.update(cls_token_loss.item(), image.shape[0])
        loss_pred_mtr.update(cls_pred_loss.item(), image.shape[0])

        optimizer.zero_grad()
        loss.backward()

        # gradient clip
        if args.clip_norm != 0:
            nn.utils.clip_grad_norm_(model.parameters(), args.clip_norm)  # clip gradient

        optimizer.step()

        batch_time.update(time.time() - end)
        end = time.time()

        writer.add_scalar('loss', loss_mtr.avg, global_step)
        writer.add_scalar('Token loss', loss_token_mtr.avg, global_step)
        writer.add_scalar('Pred loss', loss_pred_mtr.avg, global_step)

        if i % 10 == 0 or i == len(train_loader) - 1:
            progress.display(i)
        del loss


def validate(test_loader, model, args, test_stage=0):
    model.train(False)

    a_acc_list = []
    v_acc_list = []
    av_acc_list = []

    for step, (image, spec, bboxes, name) in enumerate(test_loader):
        if torch.cuda.is_available():
            spec = spec.cuda(args.gpu, non_blocking=True)
            image = image.cuda(args.gpu, non_blocking=True)
            label = bboxes['class'].cuda(args.gpu, non_blocking=True)

        a_pred_prob, v_pred_prob = model(image.float(), spec.float(), cls_target=label, mode='test')[:2]
        
        a_pred = a_pred_prob.argmax(dim=1).detach().cpu().numpy()
        v_pred = v_pred_prob.argmax(dim=1).detach().cpu().numpy()
        av_pred = (a_pred_prob * v_pred_prob).argmax(dim=1).detach().cpu().numpy()

        target = label.argmax(dim=1).detach().cpu().numpy()

        eval_index = (target < ((test_stage+1) * 9)) * (target >= ((test_stage) * 9))

        # calculate
        if eval_index.sum() != 0:
            a_acc = ((a_pred == target)*eval_index).sum() / eval_index.sum()
            v_acc = ((v_pred == target)*eval_index).sum() / eval_index.sum()
            av_acc = ((av_pred == target)*eval_index).sum() / eval_index.sum()

            a_acc_list.append(a_acc)
            v_acc_list.append(v_acc)
            av_acc_list.append(av_acc)
    
    a_acc, v_acc, av_acc = np.mean(a_acc_list), np.mean(v_acc_list), np.mean(av_acc_list)
    
    return a_acc, v_acc, av_acc


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)


class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix="", fp=None):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix
        self.fp = fp

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        msg = '\t'.join(entries)
        print(msg, flush=True)
        if self.fp is not None:
            self.fp.write(msg+'\n')

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'


if __name__ == "__main__":
    main(get_arguments())

train.py
stoneMo/CIGN
import os
import json

import torch
from torch.optim import *
import numpy as np
from sklearn import metrics
import math

from datasets import inverse_normalize
import cv2


def normalize_img(value, vmax=None, vmin=None):
    vmin = value.min() if vmin is None else vmin
    vmax = value.max() if vmax is None else vmax
    if not (vmax - vmin) == 0:
        value = (value - vmin) / (vmax - vmin)  # vmin..vmax
    return value


def visualize(raw_image, boxes):
    import cv2
    boxes_img = np.uint8(raw_image.copy())[:, :, ::-1]

    for box in boxes:

        xmin,ymin,xmax,ymax = int(box[0]),int(box[1]),int(box[2]),int(box[3])

        cv2.rectangle(boxes_img[:, :, ::-1], (xmin, ymin), (xmax, ymax), (0,0,255), 1)

    return boxes_img[:, :, ::-1]


def build_optimizer_and_scheduler_adam_v2(model, args):
    # optimizer_grouped_parameters = filter(lambda p: p.requires_grad, model.parameters())
    imgnet = []
    others = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            if 'imgnet' in name:
                imgnet.append(param)
            else:
                others.append(param)

    optimizer = Adam([{'params':imgnet}, {'params':others}], lr=args.init_lr, weight_decay=args.weight_decay)
    scheduler = None
    return optimizer, scheduler


def build_optimizer_and_scheduler_adam(model, args):
    optimizer_grouped_parameters = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = Adam(optimizer_grouped_parameters, lr=args.init_lr, weight_decay=args.weight_decay)
    scheduler = None
    return optimizer, scheduler


def build_optimizer_and_scheduler_sgd(model, args):
    optimizer_grouped_parameters = model.parameters()
    optimizer = SGD(optimizer_grouped_parameters, lr=args.init_lr)
    scheduler = None
    return optimizer, scheduler


def save_json(data, filename, save_pretty=False, sort_keys=False):
    with open(filename, mode='w', encoding='utf-8') as f:
        if save_pretty:
            f.write(json.dumps(data, indent=4, sort_keys=sort_keys))
        else:
            json.dump(data, f)


def save_results(name_list, area_list, bb_list, iou_list, conf_list, piap_list, filename):
    with open(filename, "w") as file_iou:
        file_iou.write('name,area,bb,ciou,conf\n')
        for indice in np.argsort(iou_list):
            file_iou.write(f"{name_list[indice]},{area_list[indice]},{bb_list[indice]},{iou_list[indice]},{conf_list[indice]},{piap_list[indice]}\n")


def adjust_learning_rate(optimizer, epoch, args):
    """Decay the learning rate based on schedule"""
    wu = 0 if 'warmup_epochs' not in vars(args) else args.warmup_epochs
    if args.lr_schedule == 'cos':  # cosine lr schedule
        if epoch < wu:
            lr = args.init_lr * epoch / wu
        else:
            lr = args.init_lr * 0.5 * (1. + math.cos(math.pi * (epoch - wu) / (args.epochs - wu)))
    elif args.lr_schedule == 'cte':  # constant lr
        lr = args.init_lr
    else:
        raise ValueError

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return lr
utils.py
stoneMo/CIGN