- 1. 工具
- 1. 导入模块和包
- 2. 修改系统路径 (sys.path.append)
- 3. 命令行参数解析 (argparse 模块)
- 4. 字符串格式化
- 5. 获取设备
- 6. 加载数据
- 7.设置推理结果的子路径
- 8. main() 脚本入口点
- 2. 类相关
- 1. 类的定义及初始化
- 2. 类的实例化及函数调用
- 3. 神经网络常见类
1. 工具
1. 导入模块和包
import os
import argparse
import sys as s
from accelerate import Accelerator
import
:用于导入模块和包,可以选择导入单个模块、多个模块。from ... import ...
:从特定模块中导入具体的类、函数或变量。impoert ... as ...
:可以为导入的模块指定别名,使代码简洁。
2. 修改系统路径 (sys.path.append)
# 返回当前脚本所在目录的父目录
sys.path.append(os.path.join(os.path.dirname(__file__), '..')) # 返回当前脚本所在目录的上上级目录
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
这两行代码通过 os.path.dirname(__file__)
获取当前脚本所在的目录,然后使用 ..
依次返回上一级目录。
sys.path.append()
:修改模块搜索路径,可以动态添加额外的搜索路径,以便访问其他目录中的模块。os.path.join(a, b, ...)
:将多个路径部分连接成一个完整的路径,并确保使用正确的路径分隔符(在 UNIX 系统上是/
,在 Windows 上是\
)。- 例如,
os.path.join("/home/user", "project")
将返回/home/user/project
。
- 例如,
os.path.dirname(__file__)
和..
、..
是用来构建路径的。__file__
是一个内置变量,它指向当前脚本文件的路径(字符串格式)。例如,如果脚本文件的路径是/home/user/project/scripts/train.py
,则__file__
就是这个文件的完整路径。os.path.dirname(path)
函数返回给定路径 path 的父目录(即去掉路径中的文件名部分)。例如,如果__file__
是/home/user/project/scripts/train.py
,那么os.path.dirname(__file__)
将返回/home/user/project/scripts
,即文件所在的目录。..
在路径中代表上一级目录,是一个相对路径。如果当前路径是/home/user/project/scripts
,则os.path.join("/home/user/project/scripts", '..')
会返回/home/user/project
,即当前脚本所在目录的父目录。
3. 命令行参数解析 (argparse 模块)
import argparse# 创建 ArgumentParser 对象,description 参数会显示在帮助信息中
parser = argparse.ArgumentParser(description='Train EBM model')# 添加命令行参数,指定名称、默认值、类型、帮助信息
parser.add_argument('--model_type', default="states", type=str,help='choices: states | thetas')
parser.add_argument('--dataset', default='jellyfish', type=str, help='dataset to evaluate')
parser.add_argument('--batch_size', default=4, type=int, help='Batch size for training')
parser.add_argument('--epochs', default=10, type=int, help='Number of epochs to train the model')# 解析命令行参数并将其存储在 FLAGS 对象中
FLAGS = parser.parse_args()# 使用条件判断 if-elseif-else 语句,根据 model_type 选择相应的操作
if FLAGS.model_type == "states":...
elif FLAGS.model_type == "thetas":...# 输出解析的参数值
print("Dataset:", FLAGS.dataset)
print("Batch size:", FLAGS.batch_size)
print("Epochs:", FLAGS.epochs)
在大多数深度学习训练脚本中,使用 argparse
模块来处理命令行参数。
parse = argparse.ArgumentParser()
:创建一个命令行参数解析器parse
。parser
是通过argparse.ArgumentParser()
创建的实例,它负责定义和解析命令行输入。parse.add_argument()
:添加命令行参数,指定名称、类型、默认值及帮助信息。default
:参数的默认值,如果命令行未提供该参数,则使用默认值。type
:定义参数的数据类型。help
:提供该参数的帮助描述。
parser.parse_args()
:从命令行中解析传入的参数,并将它们存储在 args 对象中。args 是一个命名空间对象,可以通过点语法访问各个参数。
运行命令行的不同示例:
- 不传递任何参数,使用默认值:
输出:# 命令行python script.py
Dataset: jellyfish Batch size: 4 Epochs: 10
- 传递自定义的命令行参数:
输出:# 命令行 python script.py --dataset "fish" --batch_size 8 --epochs 20
Dataset: fish Batch size: 8 Epochs: 20
- 查看帮助信息:
输出:python script.py --help
usage: script.py [-h] [--dataset DATASET] [--batch_size BATCH_SIZE] [--epochs EPOCHS] Train EBM model optional arguments:-h, --help show this help message and exit--dataset DATASET Dataset to evaluate--batch_size BATCH_SIZE Batch size for training--epochs EPOCHS Number of epochs to train the model
4. 字符串格式化
print("Saved at: ", results_path)
print("DATA_PATH: ", DATA_PATH)
print("number of parameters in model: ", sum(p.numel() for p in model.parameters() if p.requires_grad))
print()
:用于输出信息。- 字符串连接:
print
语句中使用逗号分隔多个变量,可以连接字符串和变量。
5. 获取设备
def get_device():return torch.device("cuda:4" if torch.cuda.is_available() else "cpu")args.device = get_device()
在深度学习任务中,通常需要指定训练的设备,如果你的机器有支持的 GPU,cuda 可以加速模型训练。如果没有 GPU,则会使用 CPU。
get_device()
这个函数会检查当前是否有可用的 GPU。
torch.device("cuda:4")
中的cuda:4
表示选择第 5 个 GPU。如果你的机器没有 5 个以上的 GPU,可能会报错。你可以使用cuda
或cuda:0
(通常默认使用第一个 GPU)。torch.cuda.is_available()
用来检查当前是否有可用的 GPU,如果没有返回 False,则会回退到 CPU。
args.device = get_device()
调用 get_device()
函数,将返回的设备对象(cuda:4
或 cpu
)赋值给 args.device
。args.device
之后可以用来指定模型的设备,确保模型训练或推理时使用正确的硬件资源。
上述代码可修改为:
def get_device():if torch.cuda.is_available():# 如果有多个 GPU,选择第一个 GPU,避免 'cuda:4' 报错return torch.device("cuda:0") else:return torch.device("cpu")args.device = get_device()
6. 加载数据
def cycle(dl):while True:for data in dl:yield data
无限循环地返回数据加载器中的数据。
7.设置推理结果的子路径
current_time = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
args.inference_result_subpath = os.path.join(args.inference_result_path,current_time + "_coeff_ratio_w_{}_J{}_".format(args.coeff_ratio_w, args.coeff_ratio_J)
)
这段代码构建推理结果保存的子目录路径 inference_result_subpath
,这个路径将包含当前时间戳以及与 coeff_ratio_w
和 coeff_ratio_J
参数相关的信息。
- 使用
datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
获取当前的时间戳,格式为年-月-日_时-分-秒
。 - 使用
args.coeff_ratio_w
和args.coeff_ratio_J
格式化字符串,这两个参数的值会被嵌入路径中。这些参数可能是用户在命令行中设置的值,影响模型训练或推理过程中的某些超参数。 - 最终路径会类似于:
inference_result_path/2024-11-17_12-30-45_coeff_ratio_w_0.3_J0.5_
。
8. main() 脚本入口点
# inference.py
def main(args):diffusion = load_model(args)dataloader = load_data(args)inference(dataloader, diffusion, args)if __name__ == '__main__':parser = argparse.ArgumentParser(description='xxx model')...args = parser.parse_args()main(args) # # 调用 main 函数
main 函数通常是脚本的核心逻辑部分,处理程序的主要任务。在 main() 函数中,可以使用 args 对象中的参数值来决定程序的行为,比如设置模型的超参数、选择训练模式等。
这段代码展示了 main 函数的核心逻辑,主要包括以下几个步骤:
-
加载模型(
diffusion = load_model(args)
):这行代码通过调用 load_model 函数加载模型及相关组件。args 参数传递了外部配置,可能包含模型路径、超参数等信息。load_model 函数返回 diffusion 扩散模型对象。 -
加载数据(
dataloader = load_data(args)
):这行代码调用 load_data 函数加载数据。args 中包含有关数据的路径或其他配置信息。dataloader 通常是一个生成数据批次(batch)的迭代器,可以用于训练或推理过程。 -
推理(
inference(dataloader, diffusion, args)
)在这行代码中,inference 函数接受输入数据的迭代器、扩散模型和配置信息参数,inference 函数通常用于执行模型的推理过程,生成预测或结果。
if __name__ == '__main__'
这一行的作用是确保当这个脚本作为主程序执行时,main() 函数会被调用。main(args)
这行代码调用 main() 函数,并将 args 对象作为参数传递给它。
因此,如果这个脚本是作为独立的 Python 文件运行的,它会执行 main 函数。
当你运行命令 python inference.py
时:
- Python 解释器会执行整个
inference.py
文件。 - 脚本中的
if __name__ == '__main__':
块会被触发。 - 然后会调用 main(args) 函数,并传入通过 argparse 解析的命令行参数 args。
这意味着,main 函数将会被执行,在 main 中使用命令行参数 args 来配置模型、加载数据、进行推理等。
2. 类相关
1. 类的定义及初始化
__init__
是 Python 中的类初始化方法,也叫构造方法。它用于在类的对象被创建时初始化对象的状态(即设置对象的属性)。
__init__
方法会在类的实例化时自动调用,并且在对象创建后执行。
class ClassName:def __init__(self, parameters):# 初始化属性或执行其他必要的操作self.attribute = value# 其他代码
__init__(self, parameters)
:__init__
方法接受至少一个参数,通常是 self,它表示类的实例对象。parameters 是该方法接受的其他参数,用于在初始化时传递值。self.attribute
:self 表示当前对象实例,可以通过它访问类的属性和方法。- value 是初始化时为属性赋的值,可以是常量、变量或通过其他逻辑生成的值。
示例 1:基础初始化方法
class Person:def __init__(self, name, age):# 初始化时将名字和年龄赋值给对象的属性self.name = nameself.age = agedef introduce(self):print(f"Hello, my name is {self.name} and I am {self.age} years old.")
在这个例子中,Person 类的 __init__
方法接受两个其他参数:name 和 age,并将它们赋值给实例对象的属性 self.name
和 self.age
。
# 创建对象时,初始化属性
person1 = Person("Alice", 30)
person2 = Person("Bob", 25)# 调用方法
person1.introduce() # 输出: Hello, my name is Alice and I am 30 years old.
person2.introduce() # 输出: Hello, my name is Bob and I am 25 years old.
示例 2:使用默认参数
class Car:def __init__(self, make, model, year=2020):# 初始化时,year 如果未传递将默认为 2020self.make = makeself.model = modelself.year = yeardef display_info(self):print(f"{self.year} {self.make} {self.model}")
在这个例子中,year 参数具有默认值 2020。如果在创建 Car 实例时未传递 year 参数,它会自动使用默认值 2020。
# 创建时传递所有参数
car1 = Car("Toyota", "Camry", 2021)# 创建时只传递 make 和 model,year 会使用默认值 2020
car2 = Car("Honda", "Civic")# 输出汽车信息
car1.display_info() # 输出: 2021 Toyota Camry
car2.display_info() # 输出: 2020 Honda Civic
2. 类的实例化及函数调用
# train_script.py# Unet3D_with_Conv3D、GaussianDiffusion、Trainer 是从模块中导入的类
from model.video_diffusion_pytorch.video_diffusion_pytorch_conv3d import Unet3D_with_Conv3D
from diffusion.diffusion_2d_jellyfish import GaussianDiffusion, Trainerif __name__ == "__main__":# 解析命令行参数并将其存储在 FLAGS 对象中FLAGS = parser.parse_args()# 创建 Unet3D_with_Conv3D 模型实例model = Unet3D_with_Conv3D(dim = 64, # 设置模型的基础维度大小out_dim = 1 if FLAGS.only_vis_pressure else 3, # 根据命令行参数 only_vis_pressure 决定输出维度dim_mults = (1, 2, 4), # 传递一个元组作为参数,用于指定每个网络层维度的倍数channels=5 if FLAGS.only_vis_pressure else 7 # 根据命令行参数 only_vis_pressure 决定通道数)# 创建 GaussianDiffusion 实例diffusion = GaussianDiffusion(model,image_size = 64,frames=FLAGS.frames,cond_steps=FLAGS.cond_steps,timesteps = 1000, # 设置扩散步骤数sampling_timesteps = 250, # 采样步骤数loss_type = 'l2', # 设置损失函数类型:L1 or L2objective = "pred_noise",device =device # 模型运行的设备(CPU/GPU))# 创建 Trainer 类的实例,该类用于管理模型的训练trainer = Trainer(diffusion,FLAGS.dataset,FLAGS.dataset_path,FLAGS.frames,FLAGS.traj_len,FLAGS.ts,FLAGS.log_path,train_batch_size = FLAGS.batch_size, # 训练的批次大小train_lr = 1e-3, # 学习率train_num_steps = 400000, # 总训练步数gradient_accumulate_every = 1, # 指定进行梯度累积的次数ema_decay = 0.995, # 用于模型参数的指数移动平均值的衰减因子save_and_sample_every = 4000, # 每 4000 步保存模型和进行采样results_path = results_path,amp = False, # 是否使用混合精度训练calculate_fid = False, # 训练过程中是否计算 fidis_testdata = FLAGS.is_testdata,only_vis_pressure = FLAGS.only_vis_pressure,model_type = FLAGS.model_type)trainer.train() # 调用 Trainer 类的 train 方法,启动模型的训练过程
这段代码展示了如何定义和使用深度学习模型的训练流程,包括模型定义、模型实例化、训练参数设置,以及如何通过面向对象编程实现模块化。if __name__ == "__main__"
使得这段代码在直接运行脚本时会执行训练逻辑,而在导入时不会执行,从而提高了代码的复用性和模块化水平。
-
if __name__ == "__main__"
:它是模块和脚本的运行入口。该语句下的代码仅在该脚本作为主程序运行时才会被执行。__name__
变量:每个 Python 模块都有一个内置属性__name__
,其值决定了模块是被导入还是直接运行。__main__
:当一个 Python 文件被直接运行时,__name__
的值会被设置为__main__
。- 导入时的行为:如果该模块被其他脚本导入,
__name__
的值是该模块的文件名(不带路径和.py
扩展名)。
-
model = ClassName(arguments)
:创建类的实例,通过类构造函数ClassName
初始化对象model
。 -
out_dim = 1 if FLAGS.only_vis_pressure else 3
:使用条件表达式(类似三元运算符)来设置输出维度,如果FLAGS.only_vis_pressure
为真,out_dim
为 1,否则为 3。
示例:
# main_script.py
if __name__ == "__main__":print("This will only run when main_script.py is executed directly.")
运行结果:
- 如果运行
python main_script.py
,将输出This will only run when main_script.py is executed directly.
- 如果
main_script.py
被其他脚本导入,如import main_script
,这行代码不会被执行。
3. 神经网络常见类
- 正弦位置编码(sinusoidal positional encoding)
常见于自然语言处理(如 Transformer)和其他序列建模任务中。正弦位置编码是一种通过正弦和余弦函数表示位置的方式,使得模型能够感知输入数据中元素的顺序。
class SinusoidalPosEmb(nn.Module):def __init__(self, dim):super().__init__()self.dim = dimdef forward(self, x):device = x.devicehalf_dim = self.dim // 2emb = math.log(10000) / (half_dim - 1)emb = torch.exp(torch.arange(half_dim, device=device) * -emb)emb = x[:, None] * emb[None, :]emb = torch.cat((emb.sin(), emb.cos()), dim=-1)return emb
Residual
类 :残差连接模块,接受一个函数作为输入, 返回该函数的输出与输入的和。
class Residual(nn.Module):def __init__(self, fn):super().__init__()self.fn = fndef forward(self, x, *args, **kwargs):return self.fn(x, *args, **kwargs) + x
Upsample
:上采样模块,使用最近邻插值方法和卷积层进行特征图的上采样。
def Upsample(dim, dim_out=None):return nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'),nn.Conv2d(dim, default(dim_out, dim), 3, padding=1))
Downsample
:下采样模块,使用Rearrange
层对特征图进行降维。
def Downsample(dim, dim_out=None):return nn.Sequential(Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1=2, p2=2),nn.Conv2d(dim * 4, default(dim_out, dim), 1))
WeightStandardizedConv2d
:实现了加权标准化卷积层,使用加权标准化(weight standardization)来提高卷积层的训练效率。
class WeightStandardizedConv2d(nn.Conv2d):def forward(self, x):eps = 1e-5 if x.dtype == torch.float32 else 1e-3weight = self.weightmean = reduce(weight, 'o ... -> o 1 1 1', 'mean')var = reduce(weight, 'o ... -> o 1 1 1', partial(torch.var, unbiased=False))normalized_weight = (weight - mean) * (var + eps).rsqrt()return F.conv2d(x, normalized_weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
LayerNorm
:自定义的层归一化,标准化每个通道的特征图。
class LayerNorm(nn.Module):def __init__(self, dim):super().__init__()self.g = nn.Parameter(torch.ones(1, dim, 1, 1))def forward(self, x):eps = 1e-5 if x.dtype == torch.float32 else 1e-3var = torch.var(x, dim=1, unbiased=False, keepdim=True)mean = torch.mean(x, dim=1, keepdim=True)return (x - mean) * (var + eps).rsqrt() * self.g
PreNorm
:先进行归一化,再传递给后续函数处理。
class PreNorm(nn.Module):def __init__(self, dim, fn):super().__init__()self.fn = fnself.norm = LayerNorm(dim)def forward(self, x):x = self.norm(x)return self.fn(x)
位置编码模块
SinusoidalPosEmb
:正弦位置编码,用于将序列的索引映射为正弦波和余弦波。
class SinusoidalPosEmb(nn.Module):def __init__(self, dim):super().__init__()self.dim = dimdef forward(self, x):device = x.devicehalf_dim = self.dim // 2emb = math.log(10000) / (half_dim - 1)emb = torch.exp(torch.arange(half_dim, device=device) * -emb)emb = x[:, None] * emb[None, :]emb = torch.cat((emb.sin(), emb.cos()), dim=-1)return emb
RandomOrLearnedSinusoidalPosEmb
:可以选择使用随机的或学习的正弦位置编码。
class RandomOrLearnedSinusoidalPosEmb(nn.Module):def __init__(self, dim, is_random=False):super().__init__()assert (dim % 2) == 0half_dim = dim // 2self.weights = nn.Parameter(torch.randn(half_dim), requires_grad=not is_random)def forward(self, x):x = rearrange(x, 'b -> b 1')freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * math.pi