Sat- nerf深度损失

首先损失函数定义在metrics.py,代码如下:

class DepthLoss(torch.nn.Module):def __init__(self, lambda_ds=1.0):super().__init__()# 初始化lambda_ds参数，用于调节深度损失的权重，并且将其缩小为原来的1/3self.lambda_ds = lambda_ds / 3.# 初始化均方误差损失函数（MSELoss），并设置reduce=False表示不对损失值进行平均self.loss = torch.nn.MSELoss(reduce=False)def forward(self, inputs, targets, weights=1.):# 创建一个字典，用来存储不同类型的深度损失loss_dict = {}# 默认使用'coarse'（粗略）类型typ = 'coarse'# 计算输入深度与目标深度之间的损失，并将结果存入字典loss_dict[f'{typ}_ds'] = self.loss(inputs['depth_coarse'], targets)# 如果输入中包含'fine'（精细）类型的深度数据if 'depth_fine' in inputs:typ = 'fine'# 计算精细深度的损失，并存入字典loss_dict[f'{typ}_ds'] = self.loss(inputs['depth_fine'], targets)# 对每个损失项应用权重for k in loss_dict.keys():# 计算加权的平均损失，并乘以lambda_ds来调整损失的权重loss_dict[k] = self.lambda_ds * torch.mean(weights * loss_dict[k])# 计算所有损失项的总和loss = sum(l for l in loss_dict.values())# 返回总损失以及包含各个深度损失的字典return loss, loss_dict

需要三个输入inputs, targets, weights=1（inputs为输入深度，target为gt，weight为权重），得到两个输出loss, loss_dict（loss为总和，loss_dict为记录单个损失的字典）。

当main函数运行到

system = NeRF_pl(args)  # 初始化 NeRF 模型系统，传入配置参数，为模型训练做好准备工作，确保所有需要的配置和资源都已经到位。

开始调用Nerf_pl 的_init_，会在其中实例化 DepthLoss 类：

self.depth_loss = DepthLoss(lambda_ds=args.ds_lambda)  # 初始化深度损失对象，传入深度监督系数

当运行trainer.fit(system）时，训练启动:

当执行trainer.fit(system)时，Lightning接管了训练过程，
Lightning首先调用prepare_data()准备数据集
然后调用configure_optimizers()设置优化器和学习率调度器

训练循环:

Lightning自动开始训练循环，每个epoch包含:

训练步骤: Lightning自动从train_dataloader()加载数据：

    def train_dataloader(self):"""创建并返回训练数据加载器字典根据配置参数创建不同模态（颜色/深度）的训练数据加载器。当self.depth为True时，会同时创建颜色数据和深度数据的加载器。数据加载器使用4个工作进程进行数据加载，启用内存锁页(pin_memory)以加速GPU数据传输，并自动进行批次数据打乱。Returns:dict: 包含数据加载器的字典，键为模态名称（"color"/"depth"），值为对应的torch.utils.data.DataLoader实例"""# 创建颜色数据的训练集加载器（第一个数据集）a = DataLoader(self.train_dataset[0],shuffle=True,num_workers=4,batch_size=self.args.batch_size,pin_memory=True)loaders = {"color": a}# 当需要加载深度数据时，创建第二个数据加载器if self.depth:b = DataLoader(self.train_dataset[1],#数据从上面dataloade 的self.train_dataset[1]，这是一个SatelliteDataset_depth类的实例，在prepare_data()方法中创建shuffle=True,num_workers=4,batch_size=self.args.batch_size,pin_memory=True)loaders["depth"] = b#通过batch["depth"]访问return loaders

可以看到深度数据从上面prepare_data()方法中创建的self.train_dataset[1]，这是一个SatelliteDataset_depth类的实例。
接下来，调用training_step()处理每个批次：

    def training_step(self, batch, batch_nb):self.log("lr", train_utils.get_learning_rate(self.optimizer))self.train_steps += 1rays = batch["color"]["rays"] # (B, 11)rgbs = batch["color"]["rgbs"] # (B, 3)ts = None if not self.use_ts else batch["color"]["ts"].squeeze() # (B, 1)results = self(rays, ts)if 'beta_coarse' in results and self.get_current_epoch(self.train_steps) < 2:loss, loss_dict = self.loss_without_beta(results, rgbs)else:loss, loss_dict = self.loss(results, rgbs)self.args.noise_std *= 0.9if self.depth:tmp = self(batch["depth"]["rays"], batch["depth"]["ts"].squeeze())kp_depths = torch.flatten(batch["depth"]["depths"][:, 0])kp_weights = 1. if self.args.ds_noweights else torch.flatten(batch["depth"]["depths"][:, 1])loss_depth, tmp = self.depth_loss(tmp, kp_depths, kp_weights)#tmp是作为imput输入进去了，kp_depths是target，kp_weights是权重if self.train_steps < self.ds_drop :loss += loss_depthfor k in tmp.keys():loss_dict[k] = tmp[k]self.log("train/loss", loss)typ = "fine" if "rgb_fine" in results else "coarse"with torch.no_grad():psnr_ = metrics.psnr(results[f"rgb_{typ}"], rgbs)self.log("train/psnr", psnr_)for k in loss_dict.keys():self.log("train/{}".format(k), loss_dict[k])self.log('train_psnr', psnr_, on_step=True, on_epoch=True, prog_bar=True)return {'loss': loss}

而在这中间 results = self(rays, ts)相当于就是隐式调用了Nerf_pl的forward算法，另外我们可以看到我们的目标self.depth_loss被调用了，tmp是作为imput输入进去了，kp_depths是target，kp_weights是权重。
那接下来我们一个一个分析这三个输入都是从哪来的。

预测深度(tmp)的完整产生过程

1. 初始数据：深度射线信息

# 在training_step中获取的初始数据，来自于train_dataloader(),
rays = batch["depth"]["rays"]  # 形状为[B, 11]的张量
ts = batch["depth"]["ts"].squeeze()  # 时间戳信息

这里的rays包含11个通道的信息：

• rays[:, 0:3]: 射线原点坐标
• rays[:, 3:6]: 射线方向向量
• rays[:, 6:7]: 近平面距离
• rays[:, 7:8]: 远平面距离
• rays[:, 8:11]: 太阳方向向量

2. 模型前向传递：self方法调用

# training_step中的调用
tmp = self(batch["depth"]["rays"], batch["depth"]["ts"].squeeze())

这调用了NeRF_pl的forward方法(在main.py中)：

def forward(self, rays, ts):chunk_size = self.args.chunkbatch_size = rays.shape[0]results = defaultdict(list)for i in range(0, batch_size, chunk_size):rendered_ray_chunks = render_rays(self.models, self.args, rays[i:i + chunk_size],ts[i:i + chunk_size] if ts is not None else None)for k, v in rendered_ray_chunks.items():results[k] += [v]for k, v in results.items():results[k] = torch.cat(v, 0)return results

该方法将射线分成小块，处理后合并结果。

3. 在forward中调用了render_rays函数：渲染逻辑

进入rendering.py中的render_rays函数：

def render_rays(models, args, rays, ts):# 获取配置参数N_samples = args.n_samples  N_importance = args.n_importancevariant = args.model  # "sat-nerf"# 分解射线信息rays_o, rays_d = rays[:, 0:3], rays[:, 3:6]near, far = rays[:, 6:7], rays[:, 7:8]# 采样深度点z_steps = torch.linspace(0, 1, N_samples, device=rays.device)z_vals = near * (1-z_steps) + far * z_steps  # 线性采样# 添加随机扰动if perturb > 0:# 采样点添加扰动代码...# 计算3D采样点坐标xyz_coarse = rays_o.unsqueeze(1) + rays_d.unsqueeze(1) * z_vals.unsqueeze(2)# 根据模型类型调用相应的inference函数if variant == "sat-nerf":from models.satnerf import inferencesun_d = rays[:, 8:11]rays_t = models['t'](ts) if ts is not None else Noneresult = inference(models["coarse"], args, xyz_coarse, z_vals, rays_d=None, sun_d=sun_d, rays_t=rays_t)# 太阳光校正相关代码...# 组织结果result_ = {}for k in result.keys():result_[f"{k}_coarse"] = result[k]# 如果需要精细采样if N_importance > 0:# 精细采样相关代码...return result_

这个函数处理射线，生成采样点，并调用合适的模型inference函数。

4. 又在render_rays中调用satnerf的inference函数

进入models/satnerf.py的inference函数：

def inference(model, args, rays_xyz, z_vals, rays_d=None, sun_d=None, rays_t=None):N_rays = rays_xyz.shape[0]N_samples = rays_xyz.shape[1]xyz_ = rays_xyz.view(-1, 3)  # 展平为[N_rays*N_samples, 3]# 处理额外输入rays_d_ = None if rays_d is None else torch.repeat_interleave(rays_d, repeats=N_samples, dim=0)sun_d_ = None if sun_d is None else torch.repeat_interleave(sun_d, repeats=N_samples, dim=0)rays_t_ = None if rays_t is None else torch.repeat_interleave(rays_t, repeats=N_samples, dim=0)# 分块运行NeRF模型chunk = args.chunkbatch_size = xyz_.shape[0]out_chunks = []for i in range(0, batch_size, chunk):out_chunks += [model(xyz_[i:i+chunk],input_dir=None if rays_d_ is None else rays_d_[i:i + chunk],input_sun_dir=None if sun_d_ is None else sun_d_[i:i + chunk],input_t=None if rays_t_ is None else rays_t_[i:i + chunk])]out = torch.cat(out_chunks, 0)# 处理输出out = out.view(N_rays, N_samples, model.number_of_outputs)rgbs = out[..., :3]       # 颜色sigmas = out[..., 3]      # 体密度# 其他输出处理...# 计算alpha合成权重deltas = z_vals[:, 1:] - z_vals[:, :-1]delta_inf = 1e10 * torch.ones_like(deltas[:, :1])deltas = torch.cat([deltas, delta_inf], -1)noise_std = args.noise_stdnoise = torch.randn(sigmas.shape, device=sigmas.device) * noise_stdalphas = 1 - torch.exp(-deltas * torch.relu(sigmas + noise))alphas_shifted = torch.cat([torch.ones_like(alphas[:, :1]), 1 - alphas + 1e-10], -1)transparency = torch.cumprod(alphas_shifted, -1)[:, :-1]weights = alphas * transparency# 计算深度depth_final = torch.sum(weights * z_vals, -1)  # 关键：计算加权平均深度# 组织返回结果result = {'rgb': rgb_final,'depth': depth_final,  # 这是最终的预测深度'weights': weights,# 其他输出...}return result

这个函数调用实际的NeRF模型，处理采样点，并生成体密度和颜色，最后计算深度。

5.inference调用SatNeRF模型

进入satnerf.py中的SatNeRF类的forward方法：

def forward(self, input_xyz, input_dir=None, input_sun_dir=None, input_t=None, sigma_only=False):# 将输入坐标送入映射层input_xyz = self.mapping[0](input_xyz)# 通过全连接网络处理xyz_ = input_xyzfor i in range(self.layers):if i in self.skips:xyz_ = torch.cat([input_xyz, xyz_], -1)xyz_ = self.fc_net[2*i](xyz_)xyz_ = self.fc_net[2*i + 1](xyz_)# 获取共享特征shared_features = xyz_# 预测体密度(sigma)sigma = self.sigma_from_xyz(shared_features)if sigma_only:return sigma# 预测颜色和其他属性xyz_features = self.feats_from_xyz(shared_features)# RGB预测...# 太阳可见度和天空颜色预测...# 预测不确定度参数betainput_for_beta = torch.cat([xyz_features, input_t], -1)beta = self.beta_from_xyz(input_for_beta)# 组合所有输出out = torch.cat([rgb, sigma, sun_v, sky_color, beta], 1)return out

这个神经网络预测每个点的体密度、颜色和其他属性。

6. 从体密度到深度的转换

回到inference函数中，关键的深度计算步骤：

# 计算alpha合成权重
alphas = 1 - torch.exp(-deltas * torch.relu(sigmas + noise))
# ...处理alphas...
weights = alphas * transparency  # 体密度权重# 预测深度：沿射线的加权平均深度
depth_final = torch.sum(weights * z_vals, -1)  # [N_rays]

体密度通过alpha合成转换为权重，然后用这些权重计算加权平均深度。

7. 结果整合与返回

最终在render_rays函数中：

result_ = {}
for k in result.keys():result_[f"{k}_coarse"] = result[k]# 深度结果作为'depth_coarse'返回

这就形成了tmp中的'depth_coarse'键，即粗略深度预测。

完整数据流向总结

1. 初始数据: batch["depth"]["rays"]和batch["depth"]["ts"] (来自数据加载器)
2. NeRF_pl.forward: 将射线分块并调用render_rays
3. render_rays: 处理射线，生成采样点，调用satnerf.inference
4. satnerf.inference:
   • 处理采样点
   • 调用SatNeRF模型预测体密度
   • 将体密度转换为权重
   • 计算加权平均深度
5. SatNeRF.forward: 神经网络计算每个点的体密度
6. 结果整合: 深度值被组织为'depth_coarse'键返回

这就是tmp中预测深度的完整产生过程，从最初的射线输入到最终的深度预测值。

SatNeRF中真实深度(kp_depths)的完整产生过程

下面我详细梳理从原始数据到最终在training_step中使用的kp_depths的完整生成过程，按照数据流向逐步说明：

1. 原始数据：卫星图像和元数据

最初的原始数据包括：

• 卫星图像（RGB图像）
• JSON元数据文件，包含相机参数(RPC)和关键点信息

这些数据存储在args.root_dir和args.img_dir指定的目录中。

2. 数据集创建：load_dataset函数

在prepare_data方法中通过调用load_dataset函数加载数据：

# main.py中的prepare_data方法
def prepare_data(self):self.train_dataset = [] + load_dataset(self.args, split="train")self.val_dataset = [] + load_dataset(self.args, split="val")

load_dataset函数(在datasets/init.py中)会根据数据类型创建不同的数据集：

def load_dataset(args, split):if args.data == 'sat':ds_list = []from .satellite import SatelliteDatasetds_list.append(SatelliteDataset(args.root_dir, args.img_dir, split, img_downscale=args.img_downscale, cache_dir=args.cache_dir))if split == 'train' and args.ds_lambda > 0:from .satellite_depth import SatelliteDataset_depthds_list.append(SatelliteDataset_depth(args.root_dir, args.img_dir, split, img_downscale=args.img_downscale, cache_dir=args.cache_dir))return ds_list# ...

当args.ds_lambda > 0时，会创建SatelliteDataset_depth类的实例作为深度数据集。

3. 深度数据集初始化：SatelliteDataset_depth

在SatelliteDataset_depth的__init__方法中处理关键点数据：

# datasets/satellite_depth.py
def __init__(self, root_dir, img_dir, split='train', img_downscale=1.0, cache_dir=None):super().__init__(root_dir, img_dir, split, img_downscale, cache_dir)

4. 如果split == “train”，调用load_train_split()方法

tie_points是通过多视图几何（MVG）从多张卫星图像中三角测量得到的3D点坐标。

def load_train_split(self):with open(os.path.join(self.json_dir, "train.txt"), "r") as f:json_files = f.read().split("\n")self.json_files = [os.path.join(self.json_dir, json_p) for json_p in json_files]if os.path.exists(self.json_dir + "/pts3d.npy"):self.tie_points = np.load(self.json_dir + "/pts3d.npy")self.all_rays, self.all_depths, self.all_ids = self.load_depth_data(self.json_files, self.tie_points, verbose=True)else:raise FileNotFoundError("Could not find {}".format(self.json_dir + "/pts3d.npy"))

读取训练图像列表
加载3D关键点（tie_points）
调用load_depth_data方法处理深度数据

5.从4中调用load_depth_data处理深度数据

load_depth_data方法是生成真实深度数据的核心：

def load_depth_data(self, json_files, tie_points, verbose=False):all_rays, all_depths, all_sun_dirs, all_weights = [], [], [], []all_ids = []kp_weights = self.load_keypoint_weights_for_depth_supervision(json_files, tie_points)for t, json_p in enumerate(json_files):# 读取JSON数据d = sat_utils.read_dict_from_json(json_p)img_id = sat_utils.get_file_id(d["img"])# 获取关键点信息pts2d = np.array(d["keypoints"]["2d_coordinates"])/ self.img_downscalepts3d = np.array(tie_points[d["keypoints"]["pts3d_indices"], :])rpc = sat_utils.rescale_rpc(rpcm.RPCModel(d["rpc"], dict_format="rpcm"), 1.0 / self.img_downscale)# 生成射线cols, rows = pts2d.Tmin_alt, max_alt = float(d["min_alt"]), float(d["max_alt"])rays = get_rays(cols, rows, rpc, min_alt, max_alt)rays = self.normalize_rays(rays)all_rays += [rays]# 获取太阳方向sun_dirs = self.get_sun_dirs(float(d["sun_elevation"]), float(d["sun_azimuth"]), rays.shape[0])all_sun_dirs += [sun_dirs]# 标准化3D坐标pts3d = torch.from_numpy(pts3d).type(torch.FloatTensor)pts3d[:, 0] -= self.center[0]pts3d[:, 1] -= self.center[1]pts3d[:, 2] -= self.center[2]pts3d[:, 0] /= self.rangepts3d[:, 1] /= self.rangepts3d[:, 2] /= self.range# 计算深度值depths = torch.linalg.norm(pts3d - rays[:, :3], axis=1)all_depths += [depths[:, np.newaxis]]# 获取权重current_weights = torch.from_numpy(kp_weights[d["keypoints"]["pts3d_indices"]]).type(torch.FloatTensor)all_weights += [current_weights[:, np.newaxis]]all_ids += [t * torch.ones(rays.shape[0], 1)]# 组合所有数据all_ids = torch.cat(all_ids, 0)all_rays = torch.cat(all_rays, 0)  # (len(json_files)*h*w, 8)all_depths = torch.cat(all_depths, 0)  # (len(json_files)*h*w, 1)all_weights = torch.cat(all_weights, 0)all_depths = torch.hstack([all_depths, all_weights])  # 深度和权重合并all_sun_dirs = torch.cat(all_sun_dirs, 0)  # (len(json_files)*h*w, 3)all_rays = torch.hstack([all_rays, all_sun_dirs])  # (len(json_files)*h*w, 11)return all_rays, all_depths, all_ids

该方法处理每个JSON文件的关键点，计算射线和深度值，并整合成训练数据。

6. 在5中调用了计算关键点权重：load_keypoint_weights_for_depth_supervision方法

def load_keypoint_weights_for_depth_supervision(self, json_files, tie_points):# 初始化权重数组kp_weights = np.ones(len(tie_points))# 收集所有关键点的2D-3D对应关系all_obs = {}for json_p in json_files:with open(json_p) as f:d = json.load(f)if "keypoints" in d.keys():# 获取RPC模型rpc = rpcm.RPCModel(d["rpc"], dict_format="rpcm")# 获取2D坐标和对应的3D索引pts2d = np.array(d["keypoints"]["2d_coordinates"]) / self.img_downscalepts3d_indices = d["keypoints"]["pts3d_indices"]# 收集观察for i, idx in enumerate(pts3d_indices):if idx not in all_obs:all_obs[idx] = []all_obs[idx].append((json_p, pts2d[i], rpc))# 计算每个3D点的重投影误差for idx, obs_list in all_obs.items():if len(obs_list) >= 2:  # 至少需要2个观察# 计算重投影误差reproj_err = compute_reprojection_error(tie_points[idx], obs_list)# 根据重投影误差设置权重kp_weights[idx] = compute_keypoint_weight(reproj_err)return kp_weights

这个方法计算每个3D关键点的权重，基于其重投影误差。重投影误差越小，权重越大。

7. 数据集的__getitem__方法：准备批次数据

__getitem__方法定义了如何访问数据集中的一个样本：

def __getitem__(self, idx):# 获取训练样本if self.train:sample = {"rays": self.all_rays[idx], "depths": self.all_depths[idx], "ts": self.all_ids[idx].long()}else:# 验证集处理...return sample

对于训练集，直接返回预先计算好的射线、深度和时间戳信息。
如果把数据集想象成一本书：

__init__相当于准备整本书和目录
__len__告诉你书有多少页
__getitem__允许你翻到任意一页并读取内容

DataLoader就像是一个阅读助手，它会按照你指定的顺序（随机或顺序）一次翻几页（批次大小）给你看。

8. 数据加载器：train_dataloader方法

在main.py的train_dataloader方法中创建数据加载器：

def train_dataloader(self):a = DataLoader(self.train_dataset[0],shuffle=True,num_workers=4,batch_size=self.args.batch_size,pin_memory=True)loaders = {"color": a}if self.depth:b = DataLoader(self.train_dataset[1],shuffle=True,num_workers=4,batch_size=self.args.batch_size,pin_memory=True)loaders["depth"] = breturn loaders

这个方法创建数据加载器，将SatelliteDataset_depth的数据作为loaders["depth"]返回。

9. 训练步骤：training_step方法

最后，在training_step方法中获取真实深度：

def training_step(self, batch, batch_nb):# ...if self.depth:tmp = self(batch["depth"]["rays"], batch["depth"]["ts"].squeeze())# 获取真实深度数据kp_depths = torch.flatten(batch["depth"]["depths"][:, 0])  # 第一列是深度值kp_weights = 1. if self.args.ds_noweights else torch.flatten(batch["depth"]["depths"][:, 1])  # 第二列是权重# 计算深度损失loss_depth, tmp = self.depth_loss(tmp, kp_depths, kp_weights)# ...

batch["depth"]["depths"][:, 0]就是我们要找的真实深度值kp_depths。

完整数据流向总结

初始数据：3D关键点坐标，存储在/pts3d.npy文件中
数据加载：load_train_split方法加载关键点并调用load_depth_data
深度计算：

从JSON文件读取2D对应关系和相机参数
计算相机射线
标准化3D点坐标
计算3D点到射线原点的距离作为深度值

权重计算：基于重投影误差计算每个点的可靠性权重
数据组织：将深度值和权重组合成all_depths张量
批次访问：通过__getitem__方法访问预计算的深度数据
训练使用：在training_step中，kp_depths从batch[“depth”][“depths”][:, 0]提取

回到training_step()

经过

        if self.depth:tmp = self(batch["depth"]["rays"], batch["depth"]["ts"].squeeze())kp_depths = torch.flatten(batch["depth"]["depths"][:, 0])kp_weights = 1. if self.args.ds_noweights else torch.flatten(batch["depth"]["depths"][:, 1])loss_depth, tmp = self.depth_loss(tmp, kp_depths, kp_weights)

就已经拿到了最后的损失了。

验证步骤: 定期调用val_dataloader()加载验证数据，并执行validation_step()

前向传播:

在training_step()和validation_step()中，代码调用了self(rays, ts)
这实际上是隐式调用了forward()方法，因为在Python中，当一个类实例被当作函数调用时，会自动调用其__call__方法，而Lightning模型的__call__会调用forward()

钩子函数:

Lightning通过一系列"钩子函数"(如training_step, validation_step等)自动组织训练流程
只需实现这些钩子函数，而不需要手动调用它们