QY的个人博客
0%

DCTdiff代码

发表于 更新于 分类于 517 2 分钟
记录分析DCTdiff代码的过程

batch = tree_map(lambda x: x.to(device), next(data_generator))# 获取一个批次的数据
为什么这里batch.Size([256,64,24])?

config里面这些参数是怎么设计的？有什么用？它和缩放参数很像，能用在缩放上面吗？另外缩放之后Y_bound怎么设计？

1
2
3
4
Y_bound=[242.382],  # eta
Y_std=[6.471, 3.588, 3.767, 2.411],  # Entropy-Based Frequency Reweighting (EBFR)
Cb_std=[4.308, 1.315, 1.487, 1.0],
Cr_std=[4.014, 1.284, 1.435, 1.0],

32*32的图像

tokens=64

block_sz=2

图像被分成多个大小为2*2的小块

self.index 是干什么的？

图像处理流程

32

1
2
3
4
5
6
7
8
9
10
11
R = img[:, :, 0]
G = img[:, :, 2]
B = img[:, :, 2]

img_y = 0.299 * R + 0.587 * G + 0.114 * B
img_cb = -0.168736 * R - 0.331264 * G + 0.5 * B + 128
img_cr = 0.5 * R - 0.418688 * G - 0.081312 * B + 128

cb_downsampled = cv2.resize(img_cb, (img_cb.shape[1] // 2, img_cb.shape[0] // 2), interpolation=cv2.INTER_LINEAR)
cr_downsampled = cv2.resize(img_cr, (img_cr.shape[1] // 2, img_cr.shape[0] // 2), interpolation=cv2.INTER_LINEAR)

  • 32*32的图
  • 分别提取R, G, B, 尺寸都是32*32
  • 使用公式转到YCbCr颜色空间
  • cb，cr降采样，从32X32降到16X16
1
2
3
4
5
6
7
8
y_blocks = split_into_blocks(img_y, self.block_sz)  # Y component, (64, 64) --> (256, 4, 4)
cb_blocks = split_into_blocks(cb_downsampled, self.block_sz)  # Cb component, (32, 32) --> (64, 4, 4)
cr_blocks = split_into_blocks(cr_downsampled, self.block_sz)  # Cr component, (32, 32) --> (64, 4, 4)

# Step 3: Apply DCT on each block
dct_y_blocks = dct_transform(y_blocks)  # (256, 4, 4) (256, 2, 2)
dct_cb_blocks = dct_transform(cb_blocks)  # (64, 4, 4) (64, 2, 2)
dct_cr_blocks = dct_transform(cr_blocks)  # (64, 4, 4) (64, 2, 2)
  • 将图像分成若干个小块(Y256, cb和cr64)，每个小块大小为2*2

函数调用

train.py 使用datasets.py中的get_dataset函数拿到数据集，然后通过DataKiader初始化数据集迭代器

分析get_dataset:

  1. 根据不同数据集的名称返回不同的类，类中有个self.train的属性，存储了数据集实例

分析程序流程

仔细！别急！

  1. 加载配置文件
  2. 初始化设置
    1. 多进程
    2. 种子
    3. 混合精读
    4. log
  3. 加载数据
    1. get_dataset加载数据集
    2. get_split根据是训练集还是测试集、有标签还是无标签切割数据集
    3. DataLoader生成一个迭代器
  4. 初始化训练状态、accelerator
    1. 将数据加载器和模型与加速器结合，以进行分布式训练。
  5. 损失重加权变量的设置
    1. train.py 74行