ai 1.1.1__py3-none-any.whl

ai/__init__.py

@@ -0,0 +1 @@
+__version__ = '1.1.1'

ai/config.py

@@ -0,0 +1,122 @@
+import os
+import yaml
+from configparser import ConfigParser
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def get_yaml_config(filepath):
+    """使用 .yaml 文件配置
+    """
+    return AttrDict(yaml.load(open(filepath, 'r'), Loader=yaml.FullLoader))
+
+
+def get_conf_config(filepath):
+    """使用 .conf 文件配置
+    """
+    conf = ConfigParser()
+    conf.read(filepath, encoding="UTF-8")
+    return conf
+
+
+class Dict(dict):
+    """重写 dict, 支持 “.” 方式的属性调用
+    """
+    def __init__(self, *args, **kwargs):
+        super(Dict, self).__init__(*args, **kwargs)
+
+    def __getattr__(self, key):
+        try:
+            value = self[key]
+            if isinstance(value, dict):
+                value = Dict(value)
+            return value
+        except KeyError as k:
+            return None
+
+    def __setattr__(self, key, value):
+        if isinstance(value, dict):
+            value = Dict(value)
+        self[key] = value
+
+    def __delattr__(self, key):
+        try:
+            del self[key]
+        except KeyError as k:
+            return None
+
+    def __call__(self, key):
+        try:
+            return self[key]
+        except KeyError as k:
+            return None
+
+
+class Config(object):
+    """支持链式调用的 Config
+    """
+    def __init__(self, filepath=None):
+        if filepath:
+            self.filepath = filepath
+        else:
+            cur_dir = os.path.split(os.path.realpath(__file__))[0]
+            self.filepath = os.path.join(cur_dir, "self.conf")
+        self.config = get_conf_config(self.filepath)
+        self.d = Dict()
+        for s in self.config.sections():
+            value = Dict()
+            for k in self.config.options(s):
+                value[k] = self.config.get(s, k)
+            self.d[s] = value
+
+    def add(self, section):
+        self.config.add_section(section)
+        self.d[section] = Dict()
+        with open(self.filepath, 'w', encoding="UTF-8") as f:
+            self.config.write(f)
+
+    def set(self, section, key, value):
+        self.config.set(section, key, value)
+        self.d[section][key] = value
+        with open(self.filepath, 'w', encoding="UTF-8") as f:
+            self.config.write(f)
+
+    def get(self, section, key):
+        return self.config.get(section, key, default=None)
+
+    def remove_section(self, section):
+        self.config.remove_section(section)
+        del self.d[section]
+        with open(self.filepath, 'w', encoding="UTF-8") as f:
+            self.config.write(f)
+
+    def remove_option(self, section, key):
+        self.config.remove_option(section, key)
+        del self.d[section][key]
+        with open(self.filepath, 'w', encoding="UTF-8") as f:
+            self.config.write(f)
+
+    def save(self):
+        for s in self.d:
+            if s not in self.config.sections():
+                self.add(s)
+            for k in self.d[s]:
+                try:
+                    v = self.get(s, k)
+                except:
+                    v = None
+                if self.d[s][k] != v:
+                    self.set(s, k, self.d[s][k])
+
+    def __getattr__(self, name):
+        if name not in self.__dict__:
+            try:
+                return self.d[name]
+            except KeyError as k:
+                self.d[name] = Dict()
+                return self.d[name]
+        return self.__dict__[name]

ai/data/augment.py

@@ -0,0 +1,450 @@
+import cv2
+import math
+import random
+import numpy as np
+from PIL import Image, ImageFilter
+
+
+class WarpMLS:
+    def __init__(self, src, src_pts, dst_pts, dst_w, dst_h, trans_ratio=1.):
+        self.src = src
+        self.src_pts = src_pts
+        self.dst_pts = dst_pts
+        self.pt_count = len(self.dst_pts)
+        self.dst_w = dst_w
+        self.dst_h = dst_h
+        self.trans_ratio = trans_ratio
+        self.grid_size = 100
+        self.rdx = np.zeros((self.dst_h, self.dst_w))
+        self.rdy = np.zeros((self.dst_h, self.dst_w))
+
+    @staticmethod
+    def _bilinear_interp(x, y, v11, v12, v21, v22):
+        return (v11 * (1 - y) + v12 * y) * (1 - x) + (v21 * (1 - y) + v22 * y) * x
+
+    def generate(self):
+        self.calc_delta()
+        return self.gen_img()
+
+    def calc_delta(self):
+        w = np.zeros(self.pt_count, dtype=np.float32)
+
+        if self.pt_count < 2:
+            return
+
+        i = 0
+        while 1:
+            if self.dst_w <= i < self.dst_w + self.grid_size - 1:
+                i = self.dst_w - 1
+            elif i >= self.dst_w:
+                break
+
+            j = 0
+            while 1:
+                if self.dst_h <= j < self.dst_h + self.grid_size - 1:
+                    j = self.dst_h - 1
+                elif j >= self.dst_h:
+                    break
+
+                sw = 0
+                swp = np.zeros(2, dtype=np.float32)
+                swq = np.zeros(2, dtype=np.float32)
+                new_pt = np.zeros(2, dtype=np.float32)
+                cur_pt = np.array([i, j], dtype=np.float32)
+
+                k = 0
+                for k in range(self.pt_count):
+                    if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                        break
+
+                    w[k] = 1. / ((i - self.dst_pts[k][0]) * (i - self.dst_pts[k][0]) +
+                                 (j - self.dst_pts[k][1]) * (j - self.dst_pts[k][1]))
+
+                    sw += w[k]
+                    swp = swp + w[k] * np.array(self.dst_pts[k])
+                    swq = swq + w[k] * np.array(self.src_pts[k])
+
+                if k == self.pt_count - 1:
+                    pstar = 1 / sw * swp
+                    qstar = 1 / sw * swq
+
+                    miu_s = 0
+                    for k in range(self.pt_count):
+                        if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                            continue
+                        pt_i = self.dst_pts[k] - pstar
+                        miu_s += w[k] * np.sum(pt_i * pt_i)
+
+                    cur_pt -= pstar
+                    cur_pt_j = np.array([-cur_pt[1], cur_pt[0]])
+
+                    for k in range(self.pt_count):
+                        if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                            continue
+
+                        pt_i = self.dst_pts[k] - pstar
+                        pt_j = np.array([-pt_i[1], pt_i[0]])
+
+                        tmp_pt = np.zeros(2, dtype=np.float32)
+                        tmp_pt[0] = np.sum(pt_i * cur_pt) * self.src_pts[k][0] - \
+                                    np.sum(pt_j * cur_pt) * self.src_pts[k][1]
+                        tmp_pt[1] = -np.sum(pt_i * cur_pt_j) * self.src_pts[k][0] + \
+                                    np.sum(pt_j * cur_pt_j) * self.src_pts[k][1]
+                        tmp_pt *= (w[k] / miu_s)
+                        new_pt += tmp_pt
+
+                    new_pt += qstar
+                else:
+                    new_pt = self.src_pts[k]
+
+                self.rdx[j, i] = new_pt[0] - i
+                self.rdy[j, i] = new_pt[1] - j
+
+                j += self.grid_size
+            i += self.grid_size
+
+    def gen_img(self):
+        src_h, src_w = self.src.shape[:2]
+        dst = np.zeros_like(self.src, dtype=np.float32)
+
+        for i in np.arange(0, self.dst_h, self.grid_size):
+            for j in np.arange(0, self.dst_w, self.grid_size):
+                ni = i + self.grid_size
+                nj = j + self.grid_size
+                w = h = self.grid_size
+                if ni >= self.dst_h:
+                    ni = self.dst_h - 1
+                    h = ni - i + 1
+                if nj >= self.dst_w:
+                    nj = self.dst_w - 1
+                    w = nj - j + 1
+
+                di = np.reshape(np.arange(h), (-1, 1))
+                dj = np.reshape(np.arange(w), (1, -1))
+                delta_x = self._bilinear_interp(di / h, dj / w,
+                                                 self.rdx[i, j], self.rdx[i, nj],
+                                                 self.rdx[ni, j], self.rdx[ni, nj])
+                delta_y = self._bilinear_interp(di / h, dj / w,
+                                                 self.rdy[i, j], self.rdy[i, nj],
+                                                 self.rdy[ni, j], self.rdy[ni, nj])
+                nx = j + dj + delta_x * self.trans_ratio
+                ny = i + di + delta_y * self.trans_ratio
+                nx = np.clip(nx, 0, src_w - 1)
+                ny = np.clip(ny, 0, src_h - 1)
+                nxi = np.array(np.floor(nx), dtype=np.int32)
+                nyi = np.array(np.floor(ny), dtype=np.int32)
+                nxi1 = np.array(np.ceil(nx), dtype=np.int32)
+                nyi1 = np.array(np.ceil(ny), dtype=np.int32)
+
+                if len(self.src.shape) == 3:
+                    x = np.tile(np.expand_dims(ny - nyi, axis=-1), (1, 1, 3))
+                    y = np.tile(np.expand_dims(nx - nxi, axis=-1), (1, 1, 3))
+                else:
+                    x = ny - nyi
+                    y = nx - nxi
+                dst[i:i + h, j:j + w] = self._bilinear_interp(
+                    x,
+                    y,
+                    self.src[nyi, nxi],
+                    self.src[nyi, nxi1],
+                    self.src[nyi1, nxi],
+                    self.src[nyi1, nxi1]
+                )
+        dst = np.clip(dst, 0, 255)
+        dst = np.array(dst, dtype=np.uint8)
+        return dst
+
+
+def distort(src, segment=4):
+    img_h, img_w = src.shape[:2]
+    cut = img_w // segment
+    thresh = cut // 3
+    if thresh == 0:
+        return src
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([np.random.randint(thresh), np.random.randint(thresh)])
+    dst_pts.append([img_w - np.random.randint(thresh), np.random.randint(thresh)])
+    dst_pts.append([img_w - np.random.randint(thresh), img_h - np.random.randint(thresh)])
+    dst_pts.append([np.random.randint(thresh), img_h - np.random.randint(thresh)])
+
+    half_thresh = thresh * 0.5
+
+    for cut_idx in np.arange(1, segment, 1):
+        src_pts.append([cut * cut_idx, 0])
+        src_pts.append([cut * cut_idx, img_h])
+        dst_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
+                        np.random.randint(thresh) - half_thresh])
+        dst_pts.append([cut * cut_idx + np.random.randint(thresh) - half_thresh,
+                        img_h + np.random.randint(thresh) - half_thresh])
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst
+
+
+def stretch(src, segment=4):
+    img_h, img_w = src.shape[:2]
+
+    cut = img_w // segment
+    thresh = cut * 4 // 5
+    if thresh == 0:
+        return src
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([0, 0])
+    dst_pts.append([img_w, 0])
+    dst_pts.append([img_w, img_h])
+    dst_pts.append([0, img_h])
+
+    half_thresh = thresh * 0.5
+
+    for cut_idx in np.arange(1, segment, 1):
+        move = np.random.randint(thresh) - half_thresh
+        src_pts.append([cut * cut_idx, 0])
+        src_pts.append([cut * cut_idx, img_h])
+        dst_pts.append([cut * cut_idx + move, 0])
+        dst_pts.append([cut * cut_idx + move, img_h])
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst
+
+
+def perspective(src):
+    img_h, img_w = src.shape[:2]
+
+    thresh = img_h // 2
+    if thresh == 0:
+        return src
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([0, np.random.randint(thresh)])
+    dst_pts.append([img_w, np.random.randint(thresh)])
+    dst_pts.append([img_w, img_h - np.random.randint(thresh)])
+    dst_pts.append([0, img_h - np.random.randint(thresh)])
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst
+
+
+def gridmask(src, d1=4, d2=16, rotate=1, ratio=0.5, mode=0):
+    h, w, c = src.shape
+    hh = math.ceil((math.sqrt(h*h + w*w)))
+    d = np.random.randint(d1, d2)
+    l = math.ceil(d * ratio)
+    
+    mask = np.ones((hh, hh), np.float32)
+    st_h = np.random.randint(d)
+    st_w = np.random.randint(d)
+    for i in range(-1, hh//d+1):
+        s = d*i + st_h
+        t = s + l
+        s = max(min(s, hh), 0)
+        t = max(min(t, hh), 0)
+        mask[s:t,:] *= 0
+    for i in range(-1, hh//d+1):
+        s = d*i + st_w
+        t = s + l
+        s = max(min(s, hh), 0)
+        t = max(min(t, hh), 0)
+        mask[:,s:t] *= 0
+    r = np.random.randint(rotate)
+    mask = Image.fromarray(np.uint8(mask))
+    mask = mask.rotate(r)
+    mask = np.asarray(mask)
+    mask = mask[(hh-h)//2:(hh-h)//2+h, (hh-w)//2:(hh-w)//2+w]
+
+    if mode == 1:
+        mask = 1 - mask
+
+    if c == 3:
+        mask = cv2.merge([mask, mask, mask])
+
+    dst = src * mask
+    return np.uint8(dst)
+
+
+class RandomDistort(object):
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        img = distort(img)
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomStretch(object):
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        img = stretch(img)
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomPerspective(object):
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        img = perspective(img)
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomGridMask(object):
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        img = gridmask(img)
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class AddPepperNoise(object):
+    """增加椒盐噪声
+    Args:
+        snr (float): Signal Noise Rate
+        p (float): 概率值，依概率执行该操作
+    """
+    def __init__(self, snr, p=0.5):
+        assert isinstance(snr, float) or (isinstance(p, float))
+        self.snr = snr
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        h, w, c = img.shape
+        signal_pct = self.snr
+        noise_pct = (1 - self.snr)
+        mask = np.random.choice((0, 1, 2), size=(h, w, 1), p=[signal_pct, noise_pct/2., noise_pct/2.])
+        mask = np.repeat(mask, c, axis=2)
+        img[mask == 1] = 255   # 盐噪声
+        img[mask == 2] = 0     # 椒噪声
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class GaussianBlur(object):
+    """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""
+
+    def __init__(self, sigma=[.1, 2.], p=0.5):
+        self.sigma = sigma
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        sigma = random.uniform(self.sigma[0], self.sigma[1])
+        img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
+        return img
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomErasing(object):
+    """
+    https://github.com/zhunzhong07/Random-Erasing
+    Args:
+        p: The prob that the operation will be performed.
+        sl: min erasing area
+        sh: max erasing area
+        r1: min aspect ratio
+        mean: erasing value
+    """
+    def __init__(self, p=0.5, sl=0.01, sh=0.1, r1=0.2, max_attempts=10, aspect_ratio=None, mean=[0.4914, 0.4822, 0.4465]):
+        self.p = p
+        self.sl = sl
+        self.sh = sh
+        self.r1 = r1
+        self.mean = mean
+        self.max_attempts = max_attempts
+        self.aspect_ratio = aspect_ratio
+       
+    def __call__(self, img):
+        if random.random() > self.p:
+            return img
+        img = np.array(img)
+        h, w, c = img.shape
+
+        for attempt in range(self.max_attempts):
+            area = h * w
+       
+            target_area = random.uniform(self.sl, self.sh) * area
+            if self.aspect_ratio is None:
+                aspect_ratio = random.uniform(self.r1, 1/self.r1)
+            else:
+                aspect_ratio = self.aspect_ratio
+
+            eh = int(round(math.sqrt(target_area * aspect_ratio)))
+            ew = int(round(math.sqrt(target_area / aspect_ratio)))
+
+            if ew < w and eh < h:
+                x1 = random.randint(0, h - eh)
+                y1 = random.randint(0, w - ew)
+                if c == 3:
+                    img[x1:x1+eh, y1:y1+ew, 0] = self.mean[0]
+                    img[x1:x1+eh, y1:y1+ew, 1] = self.mean[1]
+                    img[x1:x1+eh, y1:y1+ew, 2] = self.mean[2]
+                else:
+                    img[x1:x1+eh, y1:y1+ew, 0] = self.mean[0]
+                return Image.fromarray(np.uint8(img))
+
+        return Image.fromarray(np.uint8(img))
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'

ai/llm/postprocess.py

@@ -0,0 +1,125 @@
+def get_csv_from_llm(res):
+    if '<answer>' in res:
+        start = res.find('<answer>')
+        if start != -1:
+            res = res[start+len('<answer>'):]
+        end = res.find('</answer>')
+        if end != -1:
+            res = res[:end]
+
+    res = res.strip()
+    res = res.replace('\n\n', '\n')
+    start_tag, end_tag = '```csv', '```'
+    left = res.rfind(start_tag)
+    if left == -1:
+        start_tag = '```'
+        left = res.rfind('```')
+    right = res.rfind(end_tag)
+    if left != -1 and right != -1 and right > left:
+        res = res[left+len(start_tag):right].strip()
+    else:
+        res = res.strip('`')
+
+    return res
+
+
+def get_list_from_llm(res):
+    if '<answer>' in res:
+        start = res.find('<answer>')
+        if start != -1:
+            res = res[start+len('<answer>'):]
+        end = res.find('</answer>')
+        if end != -1:
+            res = res[:end]
+
+    res = res.strip()
+    res = res.replace('\n\n', '\n')
+    res = res.replace('"，', '",')
+
+    if '"' not in res and "'" not in res:
+        res = res.strip('[]')
+        res = res.replace('；', ';')
+        tmp = res.split(';')
+        res = '[' + ','.join([f'"{t}"' for t in tmp]) + ']'
+
+    start_tag, end_tag = '```python', '```'
+    left = res.find(start_tag)
+    if left == -1:
+        start_tag = '```'
+        left = res.find('```')
+    right = res.rfind(end_tag)
+    if left != -1 and right != -1 and right > left:
+        res = res[left+len(start_tag):right].strip()
+    elif left != -1:
+        res = res[left+len(start_tag):].strip()
+    else:
+        res = res.strip('`')
+
+    left = res.find('[')
+    if left != -1:
+        res = res[left:]
+    else:
+        res = '[' + res
+    
+    if not res.startswith('['):
+        res = '[' + res
+    if not res.endswith(']'):
+        if res.endswith(','):
+            res = res.strip(',')
+        if res.endswith(',"'):
+            res = res[:-2]
+        if not res.endswith('"'):
+            res += '"'
+        res += ']'
+    try:
+        data = eval(res)
+    except:
+        print('Error: get_list_from_llm')
+        print('res:', res)
+        data = []
+
+    return data
+
+
+def get_json_from_llm(res):
+    if '<answer>' in res:
+        start = res.find('<answer>')
+        if start != -1:
+            res = res[start+len('<answer>'):]
+        end = res.find('</answer>')
+        if end != -1:
+            res = res[:end]
+
+    res = res.strip()
+    res = res.replace('\\\"', '"').replace('\\n', '\n')
+    res = res.replace('\n\n', '\n').replace('\n', '').replace(' ', '')
+    res = res.replace('"，', '",').replace('""', '","')
+    res = res.replace('}{', '},{')
+
+    start_tag, end_tag = '```json', '```'
+    left = res.find(start_tag)
+    if left == -1:
+        start_tag = '```'
+        left = res.find(start_tag)
+    right = res.rfind(end_tag)
+    if left != -1 and right != -1 and right > left:
+        res = res[left+len(start_tag):right].strip()
+    elif left != -1:
+        res = res[left+len(start_tag):].strip()
+    else:
+        res = res.strip('`')
+
+    if '``````json' in res:
+        tmp = res.split('``````json')
+        res = '[' + ','.join(tmp) + ']'
+    
+    res = res.replace('"""', '"')
+
+    try:
+        data = eval(res)
+    except:
+        print('Error: get_json_from_llm')
+        print('res:', res)
+        data = {}
+
+    return data

ai/optim/__init__.py

@@ -0,0 +1,472 @@
+import math
+import numpy as np
+import torch
+from torch.optim import Optimizer
+from collections import defaultdict
+
+
+class AdamGC(Optimizer):
+    r"""Implements AdamGC algorithm with Gradient Centralization.
+
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False, correct_bias=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps, correct_bias=correct_bias,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(AdamGC, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdamGC, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                if group['weight_decay'] > 0:
+                    grad.add_(group['weight_decay'], p.data)
+
+                # GC operation for Conv layers and FC layers
+                length = len(list(p.data.size()))
+                if length > 1:
+                    grad.add_(-grad.mean(dim=tuple(range(1, length)), keepdim=True))
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+        return loss
+
+
+class AdamW(Optimizer):
+    r"""Implements AdamW algorithm.
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False, correct_bias=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps, correct_bias=correct_bias,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(AdamW, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdamW, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                if group['weight_decay'] > 0:
+                    grad.add_(group['weight_decay'], p.data)
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                # p.data.addcdiv_(-step_size, exp_avg, denom)
+                p.data.add_(-step_size, torch.mul(p.data, group['weight_decay']).addcdiv_(1, exp_avg, denom) )
+
+        return loss
+
+
+class AdamWGC(Optimizer):
+    r"""Implements AdamWGC algorithm with Gradient Centralization.
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, amsgrad=False, correct_bias=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps, correct_bias=correct_bias,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(AdamWGC, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(AdamWGC, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                # GC operation for Conv layers and FC layers
+                length = len(list(p.data.size()))
+                if length > 1:
+                    grad.add_(-grad.mean(dim=tuple(range(1, length)), keepdim=True))
+
+                state['step'] += 1
+
+                if group['weight_decay'] > 0:
+                    grad.add_(group['weight_decay'], p.data)
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                # p.data.addcdiv_(-step_size, exp_avg, denom)
+                p.data.add_(-step_size, torch.mul(p.data, group['weight_decay']).addcdiv_(1, exp_avg, denom) )
+
+        return loss
+
+
+class RangerGC(Optimizer):
+
+    def __init__(self, params, lr=1e-3, k=5, alpha=0.5, n_sma_threshhold=5, betas=(0.95,0.999), 
+                 eps=1e-8, weight_decay=0, amsgrad=True, transformer='softplus', smooth=50,
+                 grad_transformer='square', correct_bias=False):
+        if not 0.0 <= alpha <= 1.0:
+            raise ValueError(f'Invalid slow update rate: {alpha}')
+        if not 1 <= k:
+            raise ValueError(f'Invalid lookahead steps: {k}')
+        if not lr > 0:
+            raise ValueError(f'Invalid Learning Rate: {lr}')
+        if not eps > 0:
+            raise ValueError(f'Invalid eps: {eps}')
+
+        defaults = dict(lr=lr, k=k, alpha=alpha, step_counter=0, betas=betas, 
+                        n_sma_threshhold=n_sma_threshhold, eps=eps, weight_decay=weight_decay,
+                        smooth=smooth, transformer=transformer, grad_transformer=grad_transformer,
+                        amsgrad=amsgrad, correct_bias=correct_bias)
+        super().__init__(params, defaults)
+
+        # adjustable threshold
+        self.n_sma_threshhold = n_sma_threshhold   
+
+        # lookahead params
+        self.alpha = alpha
+        self.k = k 
+
+        # radam buffer for state
+        self.radam_buffer = [[None,None,None] for ind in range(10)]
+
+    def __setstate__(self, state):
+        super(RangerGC, self).__setstate__(state)
+
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        # Evaluate averages and grad, update param tensors
+        for group in self.param_groups:
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError('Ranger optimizer does not support sparse gradients')
+                
+                amsgrad = group['amsgrad']
+                smooth = group['smooth']
+                grad_transformer = group['grad_transformer']
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['exp_avg'] = torch.zeros_like(p_data_fp32)
+                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)                    
+
+                    # lookahead weight storage now in state dict 
+                    state['slow_buffer'] = torch.empty_like(p.data)
+                    state['slow_buffer'].copy_(p.data)
+
+                else:
+                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
+                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
+                                      
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']  
+                                
+
+                # compute variance mov avg
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                # compute mean moving avg
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                
+                # transformer
+                if grad_transformer == 'square':
+                    grad_tmp = grad**2
+                elif grad_transformer == 'abs':
+                    grad_tmp = grad.abs()
+
+                exp_avg_sq.mul_(beta2).add_((1 - beta2)*grad_tmp)
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denomc = max_exp_avg_sq.clone()
+                else:
+                    denomc = exp_avg_sq.clone()
+
+                if grad_transformer == 'square':
+                    denomc.sqrt_()                                      
+
+                # GC operation for Conv layers and FC layers
+                length = len(list(p.data.size()))
+                if length > 1:
+                    grad.add_(-grad.mean(dim=tuple(range(1, length)), keepdim=True))
+
+                state['step'] += 1
+
+                if group['weight_decay'] > 0:
+                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
+                      
+                step_size = group["lr"]
+                if group["correct_bias"]:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1 ** state["step"]
+                    bias_correction2 = 1.0 - beta2 ** state["step"]
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                if  group['transformer'] =='softplus':
+                    sp = torch.nn.Softplus(smooth)
+                    denomf = sp(denomc)
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denomf)  
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+                    p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom)
+
+                p.data.copy_(p_data_fp32)
+
+                # Integrated lookahead...
+                # We do it at the param level instead of group level
+                if state['step'] % group['k'] == 0:
+                    slow_p = state['slow_buffer'] # get access to slow param tensor
+                    slow_p.add_(self.alpha, p.data - slow_p)  # (fast weights - slow weights) * alpha
+                    p.data.copy_(slow_p)  # copy interpolated weights to RAdam param tensor
+
+        return loss

ai/tensor/__init__.py

@@ -0,0 +1,25 @@
+import torch
+
+
+def tensor2list(tensor):
+    '''将张量转换为列表
+    '''
+    return tensor.detach().cpu().tolist()
+
+
+def cat(tensors, dim=0):
+    """更高效的 torch.cat 版本，避免在只有一个元素时复制原张量
+    """
+    assert isinstance(tensors, (list, tuple))
+    if len(tensors) == 1:
+        return tensors[0]
+    return torch.cat(tensors, dim)
+
+
+def contiguous(tensor):
+    '''获取连续存储的张量
+    '''
+    if tensor.is_contiguous():
+        return tensor
+    else:
+        return tensor.contiguous()

ai/trainer/__init__.py

@@ -0,0 +1,18 @@
+import os
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed, n_gpu=1):
+    '''设置随机种子
+    '''
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if n_gpu > 0:
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True

ai/utils/__init__.py

@@ -0,0 +1,26 @@
+import os
+import bz2
+
+
+def ensure_dir(dir_path):
+    '''确保单个目录存在
+    '''
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+
+
+def ensure_dirs(dir_list):
+    '''确保多个目录存在
+    '''
+    for dir_path in dir_list:
+        ensure_dir(dir_path)
+
+
+def unpack_bz2(src_path):
+    '''解压 bz2 文件
+    '''
+    data = bz2.BZ2File(src_path).read()
+    dst_path = src_path[:-4]
+    with open(dst_path, 'wb') as fp:
+        fp.write(data)
+    return dst_path

ai-1.1.1.dist-info/METADATA

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: ai
+Version: 1.1.1
+Summary: Ai lib.
+Home-page: https://github.com/Decalogue/ai
+Author: Decalogue
+Author-email: 1044908508@qq.com
+License: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: Natural Language :: English
+Classifier: Natural Language :: Chinese (Simplified)
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Other OS
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.8
+Description-Content-Type: text/x-rst
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-python
+Dynamic: summary
+
+Ai lib.

ai-1.1.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+ai/__init__.py,sha256=E-6wAZjiLRVInABFxrWoqdYlO_t-zVhK8XNfNnaDyzI,22
+ai/config.py,sha256=rlgDA990_SdbNNtoPnB48Gg0gGkEyvxzTdMo6dplpmg,3571
+ai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai/data/augment.py,sha256=05zxXdwiXG1L-eoQd8Uz8f7wyUZhdf-a0w2jDldtrcA,14188
+ai/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai/llm/postprocess.py,sha256=37XykIrCleMJ674dqfBovab0_bLy5BMF0lFHuBnSNRQ,3314
+ai/optim/__init__.py,sha256=TJuUat9jw_6bkeQhQ5e5rXMvedV8fUpJj2N6rxwbeFw,21346
+ai/tensor/__init__.py,sha256=gXpXpnD1RLiaVSgKI9KEU4K0Hx6NJ6XZY8GGurqNXF0,542
+ai/trainer/__init__.py,sha256=LuN8w4K-Wu7KbhjuoSr0oki6Iap8YbQWNzLGJjGT3Do,425
+ai/utils/__init__.py,sha256=xGOyk8BzlHyDCNrZlakPfDewRPIXVJOOJd41MYhIDmo,493
+tests/__init__.py,sha256=83Ndli5J42Tcuq58g7r-XLkeebJaKlazC_Gi2fBezE4,108
+ai-1.1.1.dist-info/METADATA,sha256=mo3nRkCsvILXulrRcZwckI2EL1gX_gnqvwuK1gOPJjA,1194
+ai-1.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ai-1.1.1.dist-info/top_level.txt,sha256=8u8HzGvcQsdPD992eODUIMuDjQ_EdzUnYySgx4Tfw8o,9
+ai-1.1.1.dist-info/RECORD,,

ai-1.1.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ai-1.1.1.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+ai
+tests

tests/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function