解决PyTorch模型推理的非确定性：确保结果可复现的实践指南

import numpy as np
import torch
from torch import Tensor
from torchvision.models.detection import retinanet_resnet50_fpn_v2, RetinaNet_ResNet50_FPN_V2_Weights
import torchvision.transforms as T
import PIL
from PIL import Image
import random # 需要导入
import os     # 需要导入


class RetinaNet:
    def __init__(self, weights: RetinaNet_ResNet50_FPN_V2_Weights = RetinaNet_ResNet50_FPN_V2_Weights.COCO_V1):
        self.weights = weights
        # 加载预训练模型，确保使用预训练权重
        self.model = retinanet_resnet50_fpn_v2(
            weights=RetinaNet_ResNet50_FPN_V2_Weights.COCO_V1 # 明确指定权重
        )
        self.model.eval()
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model.to(self.device)
        self.transform = T.Compose([
            T.ToTensor(),
        ])

    def infer_on_image(self, image: PIL.Image.Image, label: str) -> Tensor:
        input_tensor = self.transform(image)
        input_tensor = input_tensor.unsqueeze(0)
        # 注意：input_tensor.to(self.device) 会返回一个新的张量，原张量不变
        # 正确做法是：input_tensor = input_tensor.to(self.device)
        input_tensor = input_tensor.to(self.device) # 确保输入张量在正确设备上

        with torch.no_grad():
            predictions = self.model(input_tensor)

        label_index = self.get_label_index(label)
        # 这里的打印输出显示了非确定性
        print('labels', predictions[0]['labels'])

        boxes = predictions[0]['boxes'][predictions[0]['labels'] == label_index]
        masks = torch.zeros((len(boxes), input_tensor.shape[1], input_tensor.shape[2]), dtype=torch.uint8)
        for i, box in enumerate(boxes.cpu().numpy()):
            x1, y1, x2, y2 = map(int, box)
            masks[i, y1:y2, x1:x2] = 1
        return masks

    def get_label_index(self,label: str) -> int:
        return self.weights.value.meta['categories'].index(label)

    def get_label(self, label_index: int) -> str:
        return self.weights.value.meta['categories'][label_index]

    @staticmethod
    def load_image(file_path: str) -> PIL.Image.Image:
        return Image.open(file_path).convert("RGB")

# if __name__ 部分需要添加确定性设置

# ... (其他导入) ...
import random
import os

if __name__ == '__main__':
    # --- 确保可复现性的设置 ---
    seed = 3407 # 选择一个固定整数作为随机种子

    # 1. 设置Python内置的随机数生成器
    random.seed(seed)

    # 2. 设置NumPy的随机数生成器
    np.random.seed(seed)

    # 3. 设置PyTorch的CPU随机数生成器
    torch.manual_seed(seed)

    # 4. 设置PyTorch的CUDA（GPU）随机数生成器
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed) # 为当前GPU设置种子
        torch.cuda.manual_seed_all(seed) # 为所有GPU设置种子（如果使用多GPU）

    # 5. 配置PyTorch后端以使用确定性算法
    # 强制cuDNN使用确定性算法，可能会牺牲一些性能
    torch.backends.cudnn.deterministic = True
    # 禁用cuDNN的自动调优，以确保每次都使用相同的算法
    torch.backends.cudnn.benchmark = False

    # 6. 设置Python哈希种子，影响某些哈希操作的随机性
    # 注意：此设置通常需要在Python解释器启动前完成，或在脚本开始时尽早设置
    os.environ['PYTHONHASHSEED'] = str(seed)
    # --- 确定性设置结束 ---

    from matplotlib import pyplot as plt

    image_path = 'person.jpg'
    # Run inference
    retinanet = RetinaNet()
    masks = retinanet.infer_on_image(
        image=retinanet.load_image(image_path),
        label='person'
    )
    # Plot image
    plt.imshow(retinanet.load_image(image_path))
    plt.show()
    # PLot mask
    for i, mask in enumerate(masks):
        mask = mask.unsqueeze(2)
        plt.title(f'mask {i}')
        plt.imshow(mask)
        plt.show()

# 假设您有一个数据集 my_dataset
# from torch.utils.data import DataLoader, Dataset
# class MyDataset(Dataset):
#     def __len__(self): return 100
#     def __getitem__(self, idx): return torch.randn(3, 224, 224), 0

# 在 DataLoader 初始化之前，创建并设置生成器
g = torch.Generator()
g.manual_seed(seed) # 使用与全局设置相同的种子

# 创建 DataLoader，并将生成器传入
# dataLoader = torch.utils.data.DataLoader(
#     my_dataset,
#     batch_size=32,
#     num_workers=4, # 如果 num_workers > 0，则此设置尤为重要
#     worker_init_fn=lambda worker_id: np.random.seed(seed + worker_id), # 为每个worker设置不同的种子
#     generator=g
# )