PaDiM: a Patch Distribution Modeling Framework for Anomaly Detection and Localization 代码记录

model.to(device)
model.eval()

# set model's intermediate outputs
outputs = []
def hook(module, input, output):
    outputs.append(output)

model.layer1[-1].register_forward_hook(hook)
model.layer2[-1].register_forward_hook(hook)
model.layer3[-1].register_forward_hook(hook)

train_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', [])])
test_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', [])])

# extract train set features
train_feature_filepath = os.path.join(args.save_path, 'temp_%s' % args.arch, 'train_%s.pkl' % class_name)
if not os.path.exists(train_feature_filepath):
    for (x, _, _) in tqdm(train_dataloader, '| feature extraction | train | %s |' % class_name):
        # model prediction
        with torch.no_grad():
            _ = model(x.to(device))
        # get intermediate layer outputs
        for k, v in zip(train_outputs.keys(), outputs):
            train_outputs[k].append(v.cpu().detach())
        # initialize hook outputs
        outputs = []
    for k, v in train_outputs.items():
        train_outputs[k] = torch.cat(v, 0)

    # Embedding concat
    embedding_vectors = train_outputs['layer1']
    for layer_name in ['layer2', 'layer3']:
        embedding_vectors = embedding_concat(embedding_vectors, train_outputs[layer_name])

这段代码的作用是提取训练集的特征，并将这些特征存储在 train_outputs 中。

首先，代码定义了两个有序字典 train_outputs 和test_outputs，它们的结构与之前解释的相同，用于存储各个层或模块的输出结果。

接下来，代码通过迭代训练数据加载器 train_dataloader 中的数据进行特征提取。在每次迭代中，通过将输入数据 x 传递给模型model 进行预测，其中使用了 torch.no_grad() 上下文管理器，表示在预测过程中不计算梯度，以节省内存和计算资源。

然后，代码使用一个循环遍历 train_outputs 字典的键值对，其中的 outputs 变量是之前注册的钩子函数 hook在前向传递过程中收集的中间结果。针对每个键值对，将相应的输出 v（即中间结果）附加到对应键的值列表中，并使用 cpu().detach()方法将其从GPU移动到CPU并且分离（detach）出来。在每次迭代中，完成了对该批次数据的处理后，将输出列表 outputs 重置为空列表。当遍历完成所有训练数据后，代码通过另一个循环遍历 train_outputs 字典的键值对，将值列表中的所有中间结果通过torch.cat() 方法进行拼接，形成一个整体的张量。

最后，代码将第一层的特征作为初始的嵌入向量 embedding_vectors，然后再遍历其他几个层（即 ‘layer2’ 和 ‘layer3’），将它们的特征也进行拼接，生成最终的嵌入向量。

这样，通过关联每个层或模块的输出结果，并将这些结果保存在 train_outputs字典中，代码成功提取了训练集的特征，并拼接成了包含多个层特征的嵌入向量。

# Embedding concat
embedding_vectors = train_outputs['layer1']
for layer_name in ['layer2', 'layer3']:
    embedding_vectors = embedding_concat(embedding_vectors, train_outputs[layer_name])

# randomly select d dimension
embedding_vectors = torch.index_select(embedding_vectors, 1, idx)
# calculate multivariate Gaussian distribution
B, C, H, W = embedding_vectors.size()
embedding_vectors = embedding_vectors.view(B, C, H * W)
mean = torch.mean(embedding_vectors, dim=0).numpy()
cov = torch.zeros(C, C, H * W).numpy()
I = np.identity(C)
for i in range(H * W):
    # cov[:, :, i] = LedoitWolf().fit(embedding_vectors[:, :, i].numpy()).covariance_
    cov[:, :, i] = np.cov(embedding_vectors[:, :, i].numpy(), rowvar=False) + 0.01 * I
# save learned distribution
train_outputs = [mean, cov]
with open(train_feature_filepath, 'wb') as f:
    pickle.dump(train_outputs, f)

这部分代码是对提取的特征进行处理和存储。

首先，代码通过将 train_outputs 字典中的 ‘layer1’ 特征作为初始嵌入向量 embedding_vectors。

接下来，使用循环遍历 [‘layer2’, ‘layer3’]，将这两个层的特征拼接到 embedding_vectors
中，形成新的嵌入向量。

然后，随机选择 d 维度的特征，使用 torch.index_select() 方法按照 idx 索引从bedding_vectors中选择指定维度的特征。选取的结果保存在 embedding_vectors 中。

接下来，将 embedding_vectors 的形状调整为 (B, C, H * W)，其中 B 是批次大小，C 是通道数，H 和 W分别是特征高度和宽度。

然后计算嵌入向量的多元高斯分布。首先计算 embedding_vectors 沿着 B 维度的均值，得到均值 mean，并转换为 numpy数组。

接下来，创建一个维度为 (C, C, H * W) 的全零张量 cov，用于存储协方差矩阵。同时创建一个维度为 (C, C) 的单位矩阵I。

通过循环遍历 embedding_vectors 中的每个空间位置（H * W），计算该位置的协方差矩阵。可以选择使用LedoitWolf().fit() 方法估计协方差矩阵，也可以使用 np.cov() 方法计算协方差矩阵，并添加一个小的正则项（0.01 I）以避免奇异矩阵的问题。将计算得到的协方差矩阵保存在 cov 中.*

最后，将均值 mean 和协方差矩阵 cov 存储在 train_outputs 列表中。

同时，将 train_outputs 对应的列表存储到 train_feature_filepath 文件中，使用pickle.dump() 方法进行序列化保存。

总结来说，这段代码的主要目的是基于提取的特征，计算并保存嵌入向量的均值和协方差矩阵。