在 PyCharm 中创建 conda 解释器，命名为 build_diffusion_model （建议小写+下划线）
在 build_diffusion_model 环境中执行
python -m pip install -U pip
pip install torch torchvision # 若使用CUDA，需使用selector选择CUDA版本
pip install diffusers matplotlib
python -c "import torch; print(torch.backends.mps.is_available())"
导出当前解释器环境以便复现
conda env export --no-builds > environment_macosx.yml
编写 Cell 实现模型的训练管线 mnist_baseline.ipynb
"""
依赖自检
"""
import sys
import importlib.util
print(f"{'='*10} 执行依赖自检 {'='*10}")
required_packages = [
("torch", "PyTorch" ),
("torchvision", "TorchVision"),
("diffusers", "Diffusers"),
("matplotlib", "Matplotlib")
]
missing_packages = []
for package_name, display_name in required_packages:
if importlib.util.find_spec(package_name) is None:
missing_packages.append(package_name)
else:
module = __import__(package_name)
version = getattr(module, '__version__', '未知版本')
print(f"{display_name}: {version}")
if missing_packages:
print(f"\\n【ERROR】缺少以下依赖包: {', '.join(missing_packages)}")
sys.exit(1)
else:
try:
import torch
import torchvision
from torch import nn
from torch.nn import functional as f
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt
except Exception as e:
print(f"【ERROR】导入发生未知错误")
sys.exit(1)
print(f"{'='*10} 依赖自检完成 {'='*10}")
"""
Device 自检
Apple Silicon：优先 mps，否则 cpu
NVIDIA：优先 cuda，否则 cpu
"""
print(f"{'='*10} 执行硬件自检 {'='*10}")
try:
if torch.backends.mps.is_available():
device_name = "mps"
if torch.backends.mps.is_built():
print("【INFO】Use Apple Silicon (MPS)")
elif torch.cuda.is_available():
device_name = "cuda"
print("【INFO】Use NVIDIA")
else:
device_name = "cpu"
print("【INFO】Use CPU")
device = torch.device(device_name)
x = torch.ones(1).to(device) # 在CPU中创建张量，并将其移动至device，赋值给x
print(f"{device} 可以使用")
except Exception as e:
print(f"【ERROR】{e} 设备无法使用")
device = torch.device("cpu")
print(f"{'='*10} 硬件自检完成 {'='*10}")
"""
数据集测试
"""
# 数据集
dataset = torchvision.datasets.MNIST(
root="../data/datasets",
train=True, # 使用训练集，False为测试集
download=True, # 下载数据集
transform=torchvision.transforms.ToTensor() # 将图像转换为张量
)
# 为数据集创建数据加载器
dataset_loader = DataLoader(
dataset, # 要加载的数据对象
batch_size=16, # 每次迭代加载的样本数量
shuffle=True # 打乱数据顺序
)
# 从加载器中取出第一批数据
x, y = next(iter(dataset_loader))
print('Input shape:', x.shape)
print('Labels :', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap = 'gray') # 以单通道取出所有图像，拼接成大图并用灰度显示
