chenyaofo

import pathlib
import loguru
import dataclasses

import deepspeed
import torch4x
from deepspeed import comm as dist


import pprint

chenyaofo

import asyncio
import aiofiles

tar_filenames = [f"/home/chenyaofo/datasets/imagenet-wds/train/{i:06d}.tar" for i in range(256)]
# tar_filenames = [f"/gpfs01/home/chenyaofo/imagenet-wds/train/{i:06d}.tar" for i in range(256)]

count = 0

def async_reading():
    print("asyncio reading based on naive asyncio")

chenyaofo

FROM nvidia/cuda:11.8.0-devel-ubuntu22.04

ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 \
    PATH=/opt/conda/bin:$PATH \
    PYTHON_VERSION=3.10


RUN APT_INSTALL="apt-get install -y --no-install-recommends --no-install-suggests" && \
    GIT_CLONE="git clone --depth 10" && \
    rm -rf /etc/apt/sources.list.d/cuda.list \

chenyaofo

FROM nvidia/cuda:11.8.0-devel-ubuntu22.04

ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 \
    PATH=/opt/conda/bin:$PATH \
    TZ=Asia/Shanghai \
    PYTHON_VERSION=3.9

RUN APT_INSTALL="apt-get install -y --no-install-recommends --no-install-suggests" && \
    GIT_CLONE="git clone --depth 10" && \
    rm -rf /etc/apt/sources.list.d/cuda.list \

chenyaofo

import torch
import torchvision.transforms as transforms
from PIL import Image
from io import BytesIO
from torchdata.datapipes.iter import FileLister, FileOpener, TFRecordLoader, Mapper, Shuffler, Batcher, Collator, ShardingFilter
from torchdata.dataloader2 import adapter, DataLoader2, PrototypeMultiProcessingReadingService
from codebase.torchutils.serialization import jsonunpack

from torch.utils.data import DataLoader
def get_train_transforms():

chenyaofo

import torch
import torch.nn as nn
from torchvision.models import resnet50
import types
import copy

def tbr_bn_forward_impl(self: nn.BatchNorm2d, x: torch.Tensor):
    batch_var, batch_mean = torch.var_mean(x, dim=(0, 2, 3), keepdim=True)
    batch_std = torch.sqrt(batch_var+self.eps)

chenyaofo

import torch
import torch.nn as nn
from torchvision.models import resnet50
import types
import copy


def tbr_bn_forward_impl(self: nn.BatchNorm2d, x: torch.Tensor):
    # print("x.shape", x.shape)
    batch_mean = x.mean(dim=(0, 2, 3), keepdim=True)

chenyaofo

FROM nvcr.io/nvidia/pytorch:22.11-py3

ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 \
    TZ=Asia/Shanghai

RUN APT_INSTALL="apt-get install -y --no-install-recommends --no-install-suggests" && \
    GIT_CLONE="git clone --depth 10" && \
    mirror_url=mirrors.cloud.tencent.com && \
    sed -i "s/security.ubuntu.com/$mirror_url/" /etc/apt/sources.list && \
    sed -i "s/archive.ubuntu.com/$mirror_url/" /etc/apt/sources.list && \

chenyaofo

import torch
from typing import List, Tuple
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib import font_manager

# more details about matplotlib usage at
# https://wizardforcel.gitbooks.io/matplotlib-user-guide/content/

def color(*items):

chenyaofo

import time
import os
import warnings
import pathlib

try:
    import nvidia.dali.types as types
    import nvidia.dali.fn as fn
    from nvidia.dali.plugin.pytorch import DALIGenericIterator, LastBatchPolicy
    from nvidia.dali.pipeline import Pipeline