tjyuyao

def left_aligned_gather(input, index):
    """ take along the last dim in index, while index shape is left-aligned with input and can be broadcasted.
    """
    si = list(index.shape)
    sx = list(input.shape)[len(si):]
    so = si + sx
    ex = si + [1 for _ in sx]
    index = index.view(ex).expand(so)
    return input.gather(len(si)-1, index)

tjyuyao

import inspect
import re
import torch
from torch import Tensor


def allfinite(x: Tensor) -> bool:
    return torch.isfinite(x).all().item()

tjyuyao

def pqdm(func, data, n_jobs=2):
    # pytorch based dataloader does not block for the whole results,
    # it also accepts locally defined functions.
    # These features make it favorable compared to the pqdm library or the standard multiprocessing library.

    from torch.utils.data import DataLoader
    from tqdm import tqdm

    datalen = len(data)

tjyuyao

import torch
import torch.nn as nn

def convert_pretrained(net: nn.Module, load_path: str, save_path:str):
    """convert 3rd party pretrained weights into target module assuming the state order is same.
    
    Args:
        net_type (Type): ice-learn module type.
        path (str): saved module state_dict file.
    """

tjyuyao

from __future__ import annotations
from typing import List

import paramiko
from scp import SCPClient
import time
from dataclasses import dataclass


@dataclass

tjyuyao

import torch
import torch.nn.functional as F

size = (3, 4)
H, W = size

def xy():
    x = torch.linspace(-1, 1, W)
    y = torch.linspace(-1, 1, H)

tjyuyao

import torch
import cutex

M, N, K = 4, 4, 1
a = torch.rand((M, K), dtype=torch.float32).cuda()
b = torch.rand((K, N), dtype=torch.float32).cuda()
c = torch.empty((M, N), dtype=torch.float32).cuda()

kernels = cutex.SourceModule(r"""
__global__ void matmul(Tensor<float, 2> *a, Tensor<float, 2> *b, Tensor<float, 2> *c, int M, int N, int K) {

tjyuyao

export PRJNAME="STViT"
export PRJ_ROOT="/home/hyuyao/prj/$PRJNAME"
export SRC_ROOT="$PRJ_ROOT/src/$PRJNAME"
export DATA_ROOT="$PRJ_ROOT/data"
export UTILS_ROOT="$PRJ_ROOT/utils" # put current env.sh to $PRJ_ROOT/utils

export PATH=$PATH:$UTILS_ROOT

alias cddata="cd $DATA_ROOT"
alias cdsrc="cd $SRC_ROOT && git config fetch.prune true"

tjyuyao

import torch

def steady_dropout(x, prob=0.2):
    assert len(x.shape) == 2, "expected data shape of (batch, feature), while getting " + x.shape
    n_batch = x.shape[0]
    n_feat = x.shape[1]
    n_select = int(round(n_feat * prob))
    prob = float(n_select) / float(n_feat)
    r = torch.randn(n_batch, n_feat)
    _, i = r.sort(dim=1)

tjyuyao

set -g mouse on
# to enable mouse scroll, see https://github.com/tmux/tmux/issues/145#issuecomment-150736967
bind -n WheelUpPane if-shell -F -t = "#{mouse_any_flag}" "send-keys -M" "if -Ft= '#{pane_in_mode}' 'send-keys -M' 'copy-mode -e'"