[Geração e reconstrução 3D] SSDNeRF: geração 3D e reconstrução de NeRF de difusão de estágio único

Índice de artigos da série

Título : NeRF de difusão de estágio único: uma abordagem unificada para geração e reconstrução 3D
Artigo : https://arxiv.org/pdf/2304.06714.pdf
Tarefa : Geração 3D incondicional (como gerar carros diferentes a partir do ruído, etc.), único Veja
a organização da geração 3D : Hansheng Chen,1,* Jiatao Gu,2 Anpei Chen, Tongji, Apple, Universidade da Califórnia
Código : https://github.com/Lakonik/SSDNeRF


Resumo

  Tarefas de síntese de imagens com reconhecimento de 3D, incluindo geração de cenas e nova síntese de visualização baseada em imagens. Este artigo propõe o SSDNeRF, que usa um modelo de difusão para aprender antecedentes generalizáveis ​​para campos de radiação neural (NeRF) a partir de imagens multivisualização de diferentes objetos . Estudos anteriores usaram uma abordagem em dois estágios, contando com NeRF pré-treinado como dados reais para treinar modelos de difusão. Em contraste, o SSDNeRF, como um paradigma de treinamento ponta a ponta de estágio único, otimiza conjuntamente a decodificação automática e os modelos de difusão latente do NeRF para obter reconstrução 3D simultânea e aprendizado prévio (mesmo incluindo visualizações esparsas). Ao testar, a difusão anterior pode ser gerada diretamente incondicionalmente ou combinada com observações arbitrárias de objetos invisíveis para realizar a reconstrução NeRF. SSDNeRF mostra resultados robustos comparáveis ​​ou melhores que métodos específicos de tarefa em termos de geração incondicional e reconstrução 3D de visão única/esparsa.

I. Introdução

  Com o desenvolvimento da renderização neural e dos modelos generativos , foram desenvolvidos algoritmos únicos para geração de conteúdo 3D (como reconstrução 3D de visualização única/multivisualização e geração de conteúdo 3D), mas falta uma estrutura abrangente para conectar tecnologias multitarefa. Os Campos de Radiação Neural (NeRF) resolvem o problema de renderização inversa através do ajuste de super-resolução, mostrando resultados impressionantes na síntese de novas visualizações (mas aplicáveis ​​apenas a visualizações densas e difíceis de generalizar para observações esparsas). Em contraste, muitos métodos de reconstrução 3D de visualização esparsa [pixelNeRF, Mvsnerf, ViT NeRF] dependem de codificadores feed-forward de imagem para 3D, mas não conseguem lidar com a ambigüidade de áreas ocluídas e não conseguem gerar imagens nítidas. Em termos de geração incondicional, as redes adversárias generativas (GANs) com reconhecimento de 3D [34, 5, 18, 14] são parcialmente limitadas no uso de discriminadores de imagem única, que não podem permitir que as relações de visão cruzada sejam efetivamente aprendidas a partir de multi- ver dados .

  Em trabalhos anteriores [como NeRF baseado em pontuação, Gaudi, difusão triplana, Diffrf] , ldm semelhantes foram aplicados à geração 2D e 3D, mas geralmente requerem dois estágios de treinamento. O primeiro estágio exclui o modelo de difusão e apenas pré- treina VAE (Autoencoder variável) ou decodificador automático . No entanto, no caso do nerf de difusão, acreditamos que o treinamento em dois estágios induz padrões de ruído e artefatos no código latente (especialmente com visualizações esparsas) devido à incerteza da renderização inversa, o que impede o modelo de difusão de aprender efetivamente o potencial limpo múltiplo. Para resolver este problema,

  Este artigo propõe uma estrutura unificada SSDNeRF (Single Stage Diffusion NeRF), que usa um modelo de difusão latente tridimensional (LDM) para modelar a geração de código latente da cena anterior e aprende um anterior tridimensional generalizável a partir de imagens de visualização múltipla para lidar com várias tarefas tridimensionais (Fig. 1). O paradigma de treinamento de estágio único permite o aprendizado ponta a ponta dos pesos de Difusão e NeRF. Essa abordagem combina tendências geradas e renderizadas de forma consistente para melhorar o desempenho geral e permite o treinamento em dados de visualização esparsos. Além disso, o anterior tridimensional aprendido do modelo de difusão incondicional pode ser usado para amostragem flexível de cenas de teste de observações arbitrárias.

Insira a descrição da imagem aqui

As principais contribuições são as seguintes:

1. Propor SSDNeRF, um método unificado de geração 3D incondicional e reconstrução 3D baseada em imagem; 2. Como um novo paradigma de treinamento de estágio único, SSDNeRF aprende em conjunto modelos de reconstrução e difusão NeRF
a partir de imagens multivisualização de um grande número de objetos (mesmo Existem apenas três visualizações por cena) 3. Um esquema de amostragem de ajuste fino de orientação é proposto para usar a difusão aprendida antes de realizar a reconstrução 3D a partir de qualquer número de visualizações no momento do teste.

2. Trabalho relacionado

2.1. GANs 3D

GANs foram usados ​​com sucesso para geração tridimensional   , integrando renderização baseada em projeção no gerador . Várias representações 3D foram exploradas anteriormente (nuvens de pontos, cuboides, esferas [27] e voxels); mais recentemente NeRF e campos de recursos com renderizadores de volume [34, 18, 5] e com renderizadores de malha de superfície diferenciável [14]. Os métodos acima são todos treinados com discriminadores de imagens 2D e não podem raciocinar sobre relações de visão cruzada, o que os torna fortemente dependentes do viés do modelo de consistência 3D e não podem utilizar efetivamente dados de múltiplas visualizações para aprender geometrias complexas e diversas . Gan tridimensional é usado principalmente para geração incondicional. Embora a reconstrução 3D de imagens possa ser concluída através da inversão GAN [12], a autenticidade não pode ser garantida devido à limitada capacidade de expressão potencial, conforme afirmado no artigo [Diffrf, RenderDiffusion].

2.2 Regressão e geração condicionada à visualização

  A reconstrução 3D esparsa pode ser resolvida regredindo novas visualizações da imagem de entrada.Várias arquiteturas propostas [8, 59, 28, 61] codificam a imagem em recursos de volume que são projetados para uma visualização de alvo supervisionada por meio de renderização de volume . No entanto, não conseguem raciocinar sobre a ambiguidade e produzir conteúdos diversos e significativos, o que muitas vezes leva a resultados ambíguos . Em contraste, os modelos generativos condicionados por imagem são melhores na síntese de diferentes conteúdos : 3DiM [57] propõe gerar novas visualizações a partir de um modelo de difusão de imagem condicionado pela visualização , mas este modelo carece de viés de consistência tridimensional. [Sparsefusion, Nerdi, Nerfdiff] Extraia o anterior do modelo de difusão 2D condicionado por imagem em nerf para fortalecer as restrições 3D . Esses métodos são paralelos à nossa trajetória, pois modelam relações de visão cruzada no espaço da imagem, enquanto nosso modelo é de natureza tridimensional.

2.3 Decodificador automático e difusão NeRF

  O esquema de ajuste de cena única do NeRF pode ser generalizado para múltiplas cenas , compartilhando alguns parâmetros em todas as cenas e tratando o restante como códigos de cena separados [7]. Portanto, o NeRF multicena pode ser treinado como um decodificador automático [35], aprendendo conjuntamente
bancos de códigos e compartilhando pesos do decodificador
. Com arquitetura apropriada, os códigos de cena podem ser vistos como antecedentes gaussianos latentes, permitindo a conclusão 3D e até mesmo a geração [24, 48, 38]. No entanto, como os GANs 3D, essas latentes não são expressivas o suficiente para reconstruir fielmente objetos detalhados . [Gaudi, Dos dados à functa, Rodin, etc.] Decodificadores automáticos comuns aprimorados com antecedentes de difusão latentes. DiffRF [32] utiliza DIffusion antes de realizar a conclusão 3D. Esses métodos treinam modelos de autodecodificador e difusão de forma independente em dois estágios, o que é limitado.

2.4. NeRF como autodecodificador

  Dado um conjunto de imagens bidimensionais de uma cena e parâmetros de câmera correspondentes, o campo de luz da cena pode ser ajustado no espaço tridimensional, expresso como uma função óptica y ψ (r) (onde r é usado para parametrizar o ponto final e soma de um raio na direção do espaço mundial, ψ representa o parâmetro do modelo, y∈R 3 + representa o formato RGB da luz recebida. NeRF representa o campo de luz como radiação integrada ao longo da luz que passa pelo volume tridimensional (ver meu blog para princípios específicos: [Reconstrução 3D] Princípio NeRF + explicação do código) .

  NeRF também pode ser generalizado para configurações de vários cenários, compartilhando alguns parâmetros do modelo entre todos os cenários [7]. Dadas múltiplas cenas { yijgt , rijgt }, onde yijgt , rijgt é o j-ésimo par de pixels RGB e raios da i-ésima cena , cada código de cena pode ser otimizado minimizando a perda de renderização L2 {x i } e o parâmetro compartilhado ψ :

Insira a descrição da imagem aqui
Com esse objetivo em mente, o modelo é treinado como um decodificador automático. O código da cena {xi } pode ser interpretado como um código latente. Sob a suposição de distribuições gaussianas independentes, a função óptica pode ser vista como a forma de um decodificador:
Insira a descrição da imagem aqui

2.5 Desafios na geração e reconstrução

  Um autodecodificador com pesos de treinamento ψ pode ser gerado incondicionalmente pela decodificação de códigos latentes extraídos de anteriores gaussianos. Porém, para garantir a continuidade da geração, são necessários um espaço latente de baixa dimensão e um decodificador complexo , o que aumenta a dificuldade de reconstruir verdadeiramente qualquer visão na otimização.

2.6 Modelo de difusão potencial

Modelos de difusão latente (LDM) aprendem uma distribuição anterior p ϕ (x)   em um espaço latente com parâmetros ϕ , permitindo representações latentes mais expressivas (como grades de imagens 2D). Em termos de geração de campo neural, trabalhos anteriores [2, 32, 13 , 47] adotaram um esquema de treinamento em dois estágios: primeiro treinar um decodificador automático para obter o código latente xi de cada cena e depois treiná-lo como LDM de dados reais . LDM injeta ruído gaussiano ϵ∼N (0, I) em x i , gerando código de ruído x i (t) no passo de tempo de difusão t da função empírica de escalonamento de ruído α (t) , σ (t) : = α (t) ) x i(t) ϵ . Então, uma rede de eliminação de ruído com pesos treináveis ​​ϕ remove o ruído em x i (t) para prever um código de eliminação de ruído x ^ \hat{x}x^ eu. Esta rede é comumente treinada com uma perda de eliminação de ruído L2 simplificada:
Insira a descrição da imagem aqui

w (t) é uma função de ponderação empírica dependente do tempo, e x ^ \hat{x}x^ ϕ(xi(t), t)formula uma rede de eliminação de ruído de segmentação temporal.

  1. Amostragem incondicional/guiada

  Usando os pesos treinados ϕ , vários intérpretes (por exemplo, DDIM) amostram a partir da difusão anterior p ψ (x) e eliminam recursivamente x (t) , começando com ruído gaussiano aleatório x (t) até que um estado sem ruído seja alcançado x (0) . O processo de amostragem pode ser guiado pelo gradiente da perda de renderização em relação às observações conhecidas , permitindo a reconstrução 3D da imagem no momento do teste.

  1. Limitações do treinamento em dois estágios para tarefas 3D

  LDM usando VAEs de imagem 2D geralmente é treinado em dois estágios; ao usar o decodificador automático NeRF para treinar LDM : a obtenção de um código latente expressivo por meio de otimização baseada em renderização é subdeterminada, resultando em ruído na rede de remoção de ruído (Figura 2, canto superior esquerdo); Além disso , é muito difícil reconstruir nerfs a partir de visualizações esparsas sem antecedentes aprendidos (Fig. 2, canto inferior esquerdo), o que limita o treinamento a configurações de visualização densas.
Insira a descrição da imagem aqui


3. Método deste artigo

SSDNeRF, uma estrutura que conecta um expressivo decodificador automático NeRF de três planos com um modelo de difusão latente de três planos. A Figura 3 fornece uma visão geral do modelo.

Insira a descrição da imagem aqui

3.1 Treinamento NeRF de difusão em estágio único

   Um autodecodificador pode ser visto como um VAE usando um codificador de tabela de pesquisa em vez de um codificador de rede neural típico. Portanto, os objetivos de treinamento podem ser derivados de forma semelhante aos VAEs. Utilizando o decodificador NeRF
p ψ ({y j } | x, {r j }) e a difusão latente anterior p ϕ (x) , o objetivo do treinamento é minimizar o logaritmo negativo dos dados observados { y ij gt , r ij gt } Limite superior variacional de verossimilhança (NLL) . Este artigo obtém uma perda de treinamento simplificada ignorando a incerteza (variância) no código latente:

Insira a descrição da imagem aqui
Entre eles, o código da cena {xi } , o parâmetro anterior ϕ e o parâmetro do decodificador ψ são otimizados conjuntamente em um único estágio de treinamento. Esta perda inclui a perda de renderização L rend na Equação 1 e um termo anterior de difusão na forma de NLL. Seguindo o artigo [Treinamento de máxima probabilidade de modelos de difusão baseados em pontuação, modelagem generativa baseada em pontuação no espaço latente, etc.] , substituímos a difusão NLL pelo limite superior aproximado L diff (também conhecido como destilação fracionada) na equação (2 ) . Adicionando o fator peso da experiência, o objetivo final do treinamento é:

Insira a descrição da imagem aqui

  O treinamento de estágio único, usando os códigos de cena com restrição de perdas {xi } acima , permite que o aprendido antes de completar a parte não vista, o que é particularmente benéfico para o treinamento em dados de visualização esparsos (códigos triplanos expressivos são seriamente incertos )

Balanceamento de renderização e pesos anteriores

  A proporção de peso anterior à renderização λ renddiff é a chave para o treinamento de estágio único. Para garantir a generalização, é projetado um mecanismo de ponderação empírica , no qual a perda de difusão é normalizada pela média móvel exponencial (EMA) da norma Frobenius do código de cena, expressa como: λ diff := c diff / EMA(||x i | | 2 F ), c ​​​​diff é uma escala fixa; λ rend := c rend (1−e −0.1Nv )/N v . Os pesos de renderização são determinados pelo número de visualizações visíveis Nv : A ponderação baseada em Nv é uma calibração do decodificador pψ , evitando que a perda de renderização seja dimensionada linearmente com o número de raios

Comparação com campos neurais generativos de dois estágios

  O método anterior de dois estágios [Gaudi, Diffrf, geração de campo neural 3D usando difusão triplana] ignorou o termo anterior λ diff L diff no primeiro estágio de treinamento . Isso pode ser visto como a definição dos pesos de renderização anterior λ renddiff para o infinito, resultando em códigos de cena tendenciosos e ruidosos xi . O artigo [geração de campo neural 3D usando difusão triplana] alivia parcialmente esse problema ao impor uma regularização de variação total (TV) no código da cena triplana para forçar uma priorização suave, semelhante à restrição LDM no espaço latente (coluna do meio da Figura 2 ). Control3Diff propõe aprender um modelo de difusão condicional em dados gerados por 3D GAN pré-treinados em imagens de visualização única. Em contraste, o nosso treinamento de estágio único visa incorporar diretamente a difusão antes de promover a consistência de ponta a ponta.


3.2 Amostragem guiada por imagem e ajuste fino

  Para obter uma reconstrução NeRF rápida e generalizável, cobrindo a reconstrução de visualização única a densa de visualização múltipla, propomos realizar amostragem guiada por imagem enquanto ajusta o código de amostragem, levando em consideração os antecedentes de difusão e as probabilidades de renderização . De acordo com o método de amostragem guiada por reconstrução de [Modelos de difusão de vídeo], calcula-se o gradiente de renderização aproximado g, ou seja, um código de ruído x (t) :
Insira a descrição da imagem aqui
onde, (t)(t) ) é um relação sinal-ruído com base em (SNR) (o hiperparâmetro ω é 0,5 ou 0,25). O gradiente guiado g é combinado com a previsão de pontuação incondicional, expressa como o par de saída sem ruído x ^ \hat{x}x^ Correção:

Insira a descrição da imagem aqui
A escala de orientação é λ gd . Empregamos um amostrador de correção de previsão [52] para resolver x (0) alternando uma etapa DDIM com múltiplas etapas de correção de Langevin .

  Observamos que as orientações de reconstrução não podem impor estritamente restrições de renderização para reconstruções fiéis. Para resolver este problema, reutilizamos na Equação 4, ajustando o código de cena amostrado x enquanto congelamos os parâmetros de difusão e decodificador:
Insira a descrição da imagem aqui
onde λ'diff é o peso anterior no tempo de teste, que deve ser menor que o valor do peso de treinamento λ diff (porque o conhecimento anterior do conjunto de dados de treinamento é menos confiável quando transferido para um conjunto de dados de teste diferente) . Use Adam para otimizar o código x para ajuste fino

Comparação com métodos anteriores de ajuste fino de NeRF
  Embora o ajuste fino com perdas de renderização seja comum em métodos de regressão NeRF condicionados por visualização [8, 61], nosso método de ajuste fino difere no uso de uma perda anterior de difusão no código da cena 3D. melhora significativamente a generalização para novas visões, conforme mostrado em 5.3.

3.3 Alguns detalhes

  1. Cache de gradiente anterior

A reconstrução NeRF de três planos requer pelo menos centenas de iterações de otimização   para cada código de cena xi . Entre as perdas de estágio único na fórmula (4), a perda de difusão L diff leva mais tempo para ser verificada do que a perda de renderização NeRF nativa L rend , o que reduz a eficiência geral. Para acelerar o treinamento e o ajuste fino, introduzimos uma técnica chamada Prior Gradient Caching, onde o gradiente de retropropagação em cache x λ diff L diff é reutilizado em múltiplas etapas Adam enquanto atualiza o gradiente renderizado em cada etapa. x λ rend L rend . Permite menos difusão do que renderização. A seguir está o pseudocódigo do algoritmo primário:

  1. Parametrização e ponderação para remoção de ruído

  Modelo de remoção de ruído x ^ \hat{x}x^ ϕ(x(t),t)é implementado como uma rede U-Net em DDPM (122 milhões de parâmetros no total). Sua entrada e saída são recursos de três planos com ruído e sem ruído (canais de três planos empilhados), respectivamente. Para a forma do teste, adotamosa parametrizaçãovv^\hat{v}[43: Destilação progressiva para amostragem rápida de modelos de difusão]v^ ϕ(x(t),t), de modo quex ^ \hat{x}x^ = α(t)x(t)−σ(t)v ^ \hat{v}v^ . Em relação à função de ponderação w(t), LSGM [54] adota dois mecanismos diferentes para otimizar oxie o peso de difusãoϕ;Em contraste, observamos que a ponderação baseada na relação sinal-ruído w (t)= (α(t)/σ (t))usada na Equação 5tem um bom desempenho.


4. Experimente

41 conjuntos de dados

  Os experimentos usam conjuntos de dados ShapeNet SRN [6,48] e Amazon Berkeley Objects (ABO) Tables [9]. O conjunto de dados SRN fornece dois tipos de cenas de objeto único, ou seja, carros e cadeiras.A divisão de vagões de trem/teste é 2458/703 e a de cadeiras é 4612/1317. Cada cena de treinamento possui 50 visualizações aleatórias de uma esfera e cada cena de teste possui 251 visualizações em espiral do hemisfério superior. O conjunto de dados ABO Tables fornece uma divisão de treinamento/teste de 1520/156 cenas de tabela, cada uma com 91 visualizações do hemisfério superior. A resolução de renderização é 128×128.

4.2 Geração incondicional

  A geração incondicional é avaliada usando os conjuntos de dados SRN Cars e ABO Tables. O conjunto de dados Carros apresenta desafios na geração de texturas nítidas e complexas, enquanto o conjunto de dados Tabelas consiste em diferentes geometrias e materiais realistas. O modelo é treinado 1 milhão de vezes em todas as imagens do conjunto de treinamento.

  1. Esquemas e métricas de verificação

  Para carros SRN , seguindo Functa, amostramos 704 cenas do modelo de difusão e renderizamos cada cena usando as 251 poses de câmera fixas no conjunto de teste . Para a tabela ABO, 1000 cenas foram amostradas de acordo com DiffRF e cada cena foi renderizada usando 10 câmeras aleatórias. Gere anotações de métricas de qualidade: Frechet Inception Distance (FID) e Kernel Inception Distance (KID) .

  1. Comparação com métodos de última geração :
    Conforme mostrado na Tabela 1, em carros SRN, o SSDNeRF de um estágio supera significativamente o EG3D no KID (mais adequado para pequenos conjuntos de dados). Ao mesmo tempo, seu FID é significativamente melhor que o Functa (que usa um LDM, mas possui um código latente de baixa dimensão). Entre as tabelas ABO, o SSDNeRF tem um desempenho significativamente melhor que o EG3D e o DiffRF.

  2. Estágio único e dois estágios :
    Nos carros SRN, foram comparados o treinamento em estágio único e o treinamento regularizado TV em dois estágios sob a mesma arquitetura de modelo.Os resultados na Tabela 1 mostram as vantagens do estágio único.

Insira a descrição da imagem aqui

4.3 Reconstrução NeRF de visão esparsa

  Os dados de Carros apresentam o desafio de recuperar diferentes texturas, e os dados de Cadeira requerem a reconstrução precisa de diferentes formas. O modelo é treinado em todas as imagens do conjunto de treinamento por 8.000 iterações, e descobrimos que cronogramas mais longos levam à redução do desempenho na reconstrução de objetos invisíveis. Este comportamento é consistente com os resultados da interpolação

  1. Planos e Indicadores de Avaliação

  A métrica de avaliação do PixelNeRF [59] é usada: dada a imagem de entrada da amostra da cena de teste, o código da cena de três planos é obtido por meio de ajuste fino de orientação e a nova qualidade de visualização é avaliada para a imagem em perspectiva desconhecida. Indicadores de qualidade de imagem : taxa média de modulação de sinal de pico (PSNR), similaridade estrutural (SSIM) e similaridade de patch de imagem perceptual aprendida (LPIPS) [60]. e FID entre imagens sintéticas e reais (como 3DiM).

  1. Compare com outros métodos

  Na Tabela 2: SSDNeRF possui o melhor L PIPS e a melhor fidelidade perceptual. 3DiM produz imagens de alta qualidade (melhor FID), mas tem a menor fidelidade à verdade fundamental (PSNR); CodeNeRF relata o melhor PSNR em carros de visão única, mas sua expressividade limitada resulta em saída borrada (Figura 5) e L PIPS e FID; VisionNeRF alcança desempenho equilibrado em todas as métricas de visualização única, mas pode ter dificuldade em gerar detalhes de textura no lado invisível do carro (por exemplo, o outro lado da ambulância na Figura 5). Além disso, o SSDNeRF tem vantagens óbvias na reconstrução de visão dupla, alcançando o melhor desempenho em todas as métricas relevantes.

Insira a descrição da imagem aqui

Insira a descrição da imagem aqui

4.4 Treinamento SSDNeRF em dados de visualização esparsos

Esta seção treina SSDNeRF em um subconjunto de visualização esparsa do conjunto completo de treinamento SRN Cars, selecionando aleatoriamente um conjunto fixo de três visualizações em cada cena . Observe que haverá uma queda razoável de desempenho esperada em comparação com o treinamento de visualização densa, uma vez que todo o conjunto de dados de treinamento foi reduzido para 6% de seu tamanho original.

  1. Gerar incondicionalmente

No meio do treinamento, os códigos dos três planos são redefinidos para sua média. Isso ajuda a evitar que o modelo fique preso em um mínimo local com artefatos geométricos sobreajustados. Dobre o tempo de treinamento de acordo. O modelo alcançou um bom FID de 19,04±1,10 e um KID/10−3 de 8,28±0,60. Os resultados são mostrados na Figura 7.

Insira a descrição da imagem aqui

  1. Reconstrução de visão única

Adotamos a mesma estratégia de treinamento de 5.3. Com a nossa abordagem de ajuste fino de orientação, o modelo atinge uma pontuação LPIPS de 0,106, ainda melhor do que a maioria dos métodos anteriores usando o conjunto de treinamento completo na Tabela 2.

  1. Comparação com regularização de TV

A Figura 8 (b) mostra as imagens RGB e geometrias representadas pelos códigos latentes da cena aprendidos nas três visualizações durante o treinamento. Em contraste, decodificadores automáticos comuns de três planos com regularização de TV (Figura 8 (a)) muitas vezes não conseguem reconstruir cenas a partir de visualizações esparsas, resultando em graves artefatos geométricos. Portanto, treinar modelos de dois estágios com latência de expressão em dados de visualização esparsos era anteriormente inviável.

4.5 Interpolação NeRF

     De acordo com as configurações do DDIM, dois valores iniciais x (T) ∼N(0,I) são amostrados, interpolados usando interpolação linear esférica [46], e então o solucionador determinístico é usado para obter as amostras interpoladas. No entanto, como apontado em [37, 40], o modelo de difusão gaussiana padrão geralmente leva a uma interpolação não suave . No SSDNeRF (resultados mostrados na Figura 9), descobrimos que um modelo treinado com reconstrução de visão esparsa de parada inicial (a) produz transições razoavelmente suaves, enquanto um modelo treinado com geração incondicional (b) produz amostras distintas, mas descontínuas. Isto sugere que a parada precoce preserva uma priorização mais suave, levando a uma melhor generalização para a reconstrução de visão esparsa.

Insira a descrição da imagem aqui

5. Código

5.1 Geração incondicional

Selecione o arquivo de configuração aqui : ssdnerf_cars_uncond.py e baixe o peso correspondente : ssdnerf_cars_uncond_1m_emaonly.pth. O significado do arquivo de configuração é mostrado abaixo:

ssdnerf_cars3v_uncond
   │      │      └── testing data: test unconditional generation
   │      └── training data: train on Cars dataset, using 3 views per scene
   └── training method: single-stage diffusion nerf training

stage2_cars_recons1v
   │     │      └── testing data: test 3D reconstruction from 1 view
   │     └── training data: train on Cars dataset, using all views per scene
   └── training method: stage 2 of two-stage training

提示:权重可能不是跟模型完全匹配,猜想可能因为在重建中还会训练;另外就是想引入一些网络的随机权重,增强泛化性。此外,大部分重要函数,都用C++做了编译(便于cuda并行计算),因此无法解析

1. Gere ruído como entrada

num_batches = len(data['scene_id'])                              # 6
noise = torch.randn((num_batches, *self.code_size), 
                     device=get_module_device(self))             # (8,3,6,128,128)

code_diff = code_diff.reshape(code.size(0), *self.code_reshape)  # (8,18,128,128

2. processo de difusão:
ddim_sample da classe: class GaussianDiffusion (nn.Module):

sample_fn_name = 'ddim'

# 参数设置---------------------------------------------------------------------
x_t = noise
num_timesteps = self.test_cfg.get('num_timesteps', self.num_timesteps)  # 50
langevin_steps = self.test_cfg.get('langevin_steps', 0)                 # 0
langevin_t_range = self.test_cfg.get('langevin_t_range', [0, 1000])     # [0,1000]
timesteps = torch.arange(start=self.num_timesteps - 1, end=-1, 
                 step=-(self.num_timesteps / num_timesteps)).long().to(device)
# 从01000步中,随机选50个时间步

# 50次的逆扩散过程 (最重要!!)-------------------------------------------------
for step, t in enumerate(timesteps):                 # 0, 999
    t_prev = timesteps[step + 1]                     # 979
    x_t, x_0_pred = self.p_sample_ddim( x_t, t, t_prev, concat_cond=None, cfg=self.test_cfg, **kwargs)
    # 以上 self.p_sample_ddim过程,下面有详解

	eps_t_pred = (x_t - self.sqrt_alphas_bar[t] * x_0_pred) / self.sqrt_one_minus_alphas_bar[t]
	pred_sample_direction = np.sqrt(1 - alpha_bar_t_prev - tilde_beta_t * (eta ** 2)) * eps_t_pred
	x_prev = np.sqrt(alpha_bar_t_prev) * x_0_pred + pred_sample_direction      # (8,18,128,128)

    x_t = x_prev              # 更新x_t,再次循环

2.1 Em self.p_sample_ddim, vários parâmetros predefinidos: fator de difusão β \betab、 a 、a ˉ \bar{a}aespere _

def prepare_diffusion_vars(self):

	## 1.扩散因子 beta :-----------------------------------------------------------
	scale = 1000 / diffusion_timesteps      # 一般是(1000步)
	beta_0 = scale * 0.0001
	beta_T = scale * 0.02
	self.betas = np.linspace(beta_0, beta_T, diffusion_timesteps, dtype=np.float64)
	
	## 2.alphas及其他参数 :--------------------------------------------------------
	self.alphas = 1.0 - self.betas
	self.alphas_bar = np.cumproduct(self.alphas, axis=0)
	self.alphas_bar_prev = np.append(1.0, self.alphas_bar[:-1])
	self.alphas_bar_next = np.append(self.alphas_bar[1:], 0.0)
	
	## 3. 计算扩散概率 q(x_t | x_0)--------------------------------------------
	self.sqrt_alphas_bar = np.sqrt(self.alphas_bar)
	self.sqrt_one_minus_alphas_bar = np.sqrt(1.0-self.alphas_bar)
	self.log_one_minus_alphas_bar = np.log(1.0-self.alphas_bar)
	self.sqrt_recip_alplas_bar = np.sqrt(1.0 / self.alphas_bar)
	self.sqrt_recipm1_alphas_bar = np.sqrt(1.0 / self.alphas_bar-1)
	
	##  4.计算后验概率 q(x_{
    
    t-1} | x_t, x_0)  -----------------------------------
	self.tilde_betas_t = self.betas * (1-self.alphas_bar_prev) / (1-self.alphas_bar)
	    
	## 5.clip log var for tilde_betas_0 = 0----------------------------------------
	self.log_tilde_betas_t_clipped = np.log(np.append(self.tilde_betas_t[1], self.tilde_betas_t[1:]))
	self.tilde_mu_t_coef1 = np.sqrt(self.alphas_bar_prev) / (1 - self.alphas_bar) * self.betas
	self.tilde_mu_t_coef2 = np.sqrt(self.alphas) * (1 - self.alphas_bar_prev) / (1 - self.alphas_bar)

2.2 Processo de difusão inversa em self.p_sample_ddim: Unet padrão (com operação de autoatenção), as dimensões frontal e traseira permanecem inalteradas

	alpha_bar_t_prev = self.alphas_bar[t_prev]  
	embedding = self.time_embedding(t)    # bs*[999] --> (bs,512)
	h, hs = x_t, []
	
	## 0.如果有condition,跟噪声x_t拼接
	if self.concat_cond_channels > 0:
	    h = torch.cat([h, concat_cond], dim=1)
	    
	## 1.下采样阶段 
	for block in self.in_blocks:
	    h = block(h, embedding)
	    hs.append(h)                           # (8,512,8,8)
	
	# 2.中间阶段
	h = self.mid_blocks(h, embedding)
	
	# 上采样
	for block in self.out_blocks:
	    h = block(torch.cat([h, hs.pop()], dim=1), embedding)    #   (8,128,128,128)
	denoising_output = self.out(h)                    

return denoising_output                                # (8,18,128,128)

2.3 Em self.p_sample_ddim, obtenha o x_0 previsto

denoising_output = self.denoising(x_t,t,concat_cond=None)   # 上步结果(8,18,128,128if self.denoising_mean_mode.upper() == 'V':
    x_0_pred = sqrt_alpha_bar_t * x_t - sqrt_one_minus_alpha_bar_t * denoising_output
    # 预测的x_0 是 x_t和denoising_output结果的线性组合

if clip_denoised:
    x_0_pred = x_0_pred.clamp(*clip_range)        # 按-22 截断,防止梯度消失/爆炸
  1. self.get_density:A densidade de predição NeRF
    está localizada na classe grande BaseNeRF(nn.Module):
density_grid = self.get_init_density_grid(num_scenes=8, device)

    tmp_grid = torch.full_like(density_grid, -1)     # (8,262144*[-1]


    # 生成64*64*64的网格:--------------------------------------------------------------------------
    if iter_density < 16:
        X = torch.arange(self.grid_size, dtype=torch.int32, device=device).split(S)   # 64
        Y = torch.arange(self.grid_size, dtype=torch.int32, device=device).split(S)
        Z = torch.arange(self.grid_size, dtype=torch.int32, device=device).split(S)

        # 只循环一次:
        for xs in X:
            for ys in Y:
                for zs in Z:
                    # 构建三维网格----------------------------------------------------------------------
                    xx, yy, zz = custom_meshgrid(xs, ys, zs)         # (64,64,64)
                    coords = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)],
                                       dim=-1)                       # [262144, 3]
                    indices = morton3D(coords).long()                # (262144): 打乱的索引
                    xyzs = (coords.float() - (self.grid_size - 1) / 2) * (2 * decoder.bound / self.grid_size)  # 以立方体中心点做归一化到(-1,1)
                    # 添加噪声--------------------------------------------------------------------------
                    half_voxel_width = decoder.bound / self.grid_size    # 1/64
                    xyzs += torch.rand_like(xyzs) * (2 * half_voxel_width) - half_voxel_width   
                    # 解码出密度值----------------------------------------------------------------------
                    sigmas = decoder.point_decode(xyzs , code)       # 下面有解析
                    tmp_grid[:, indices] = sigmas.clamp(max=torch.finfo(tmp_grid.dtype).max).to(tmp_grid.dtype)
            
                    # 求有效密度------------------------------------------------------------------------
		            valid_mask = (density_grid >= 0) & (tmp_grid >= 0)                                 # (8,262144)
		            density_grid[:] = torch.where(valid_mask, torch.maximum(density_grid * decay, tmp_grid), density_grid)   # (8,262144)
		            mean_density = torch.mean(density_grid.clamp(min=0))  # -1 regions are viewed as 0 density.       # 352.75
		            iter_density += 1
		
		            # 保存结果到 bitfield ----------------------------------------------------------------
		            density_thresh = min(mean_density, density_thresh)     # min(352.75, 0.1)
		            packbits(density_grid, density_thresh, density_bitfield)           
		            # 将 grid 与 thresh 作比较,每8个字节为一组,大于的存在 bitfield 中


def point_decode(self, xyzs, dirs, code, density_only=False):
    
    num_scenes, _, n_channels, h, w = code.size()        # 8,6,128,128

    # # 坐标上插值code特征:
    point_code = F.grid_sample(code.reshape(24,6,128,128),    
                               self.xyz_transform(xyzs),      # (24,1,262144,2) .取出其中的3个面坐标,并拼接
                               mode=self.interp_mode, padding_mode='border', 
                               align_corners=False).reshape(8,3,6, 262144)
    point_code = point_code.permute(0, 3, 2, 1).reshape( 2097152, 18)                         
    
    base_x = self.base_net(point_code)                 # linear(2097152, 18)-> (2097152, 64) 
    base_x_act = self.base_activation(base_x)
    sigmas = self.density_net(base_x_act).squeeze(-1)  # (2097152, 64) -> (2097152, 1) 
    rgb = None if density_only
  1. Renderizar imagem (decodificar)
image, depth = self.render(decoder, code, density_bitfield, h, w, test_intrinsics, test_poses, cfg=cfg)      

Entrada de função: # código: (8,3,6,128,128) campo de bits: (8,32768) pose: (8,251,4,4) intrínsecos: [131.250, 131.250, 64,0, 64,0] A
seguir está a expansão específica:

## 1. 得到射线起点和方向---------------------------------------------------
rays_o, rays_d = get_cam_rays(poses, intrinsics, h, w)  #(8,251,128,128,3) 拓展部分有解析

## 2.VolumeRenderer-----------------------------------------------------
outputs = decoder(rays_o, rays_d, code, density_bitfield, grid_size, dt_gamma=0, perturb=False)
   
    # 2.1查询每条射线跟 aabb(-1,-1,-1,1,1,1)交点---------------------------
    nears, fars = batch_near_far_from_aabb(rays_o, rays_d, self.aabb, self.min_near=0.2)
    #  nears、fars:(8, 4112384=251*128*128)

    # 2.2 渲染256次                

          step = 0
          while step < self.max_steps:
                # 射线上均匀采样空间点,插值得到对应密度 delta--------------------
                xyzs, dirs, deltas = march_rays(4112384, n_step, rays_alive, rays_t,
                                     rays_o, rays_d_single, self.bound=1, bitfield, 1, 
                                     grid_size, nears, fars,  align=128, perturb=False, dt_gamma=0, max_steps=256)

                sigmas, rgbs,= self.point_decode(xyzs, dirs, code) # xyzs(4112512, 3)  code(3,6,128,128)
                   
                    # point_decode展开如下:---------------------------------------
	                point_code_single = F.grid_sample(code_single, self.xyz_transform(xyzs_single),
	                                         mode='bilinear', padding_mode='border', align_corners=False ).squeeze(-2) # (3,6,4112512)


			        base_x = self.base_net(point_code)       # linear (4112512, 18) -> (4112512, 64) 
			        base_x_act = self.base_activation(base_x)
			        sigmas = self.density_net(base_x_act).squeeze(-1)     # (4112512, 64) -> (4112512, 1) 

		            # 渲染颜色(因为颜色需要位置信息+方向信息)----------------------------
		            sh_enc = sh_encode(dirs, self.degree=4)   # 三维坐标,位置编码至16维
		            color_in = self.base_activation(base_x + self.dir_net(sh_enc))  # linear+Silu->(4112512,64)
		            rgbs = self.color_net(color_in)    # linear:(64,3)
		            if self.sigmoid_saturation > 0:
		                rgbs = rgbs * (1 + self.sigmoid_saturation * 2) - self.sigmoid_saturation

                # 光线追踪和 图像生成:
                composite_rays(4112512, rays_t, sigmas, rgbs, deltas, depth, image, T_thresh=0.001)

A função composite_rays realiza cálculos iterativos em cada raio e, finalmente, obtém as informações da imagem composta com base nas etapas de propagação da luz, acumulação de valores de cor e atualização de parâmetros.

Por fim, são obtidos três resultados (batchsize=8):imagem(4112384,3),soma_pesos(4112384),profundidade(4112384).Após a operação de remodelação, a imagem final é obtida:

weights = torch.stack(outputs['weights_sum'], dim=0) 
rgbs = (torch.stack(outputs['image'], dim=0) + self.bg_color * (1 - weights.unsqueeze(-1))     # self.bg_color = 1
depth = torch.stack(outputs['depth'], dim=0) 

5.2 Treinamento de reconstrução NeRF (multivisualização): ssdnerf_cars_recons1v

O arquivo de configuração é: ssdnerf_cars_recons1v.py. Selecione 50 perspectivas de cada vez para reconstrução.

1. Perda do processo de difusão (o code_list_ em cache é muito importante):

# 1.加载缓存cache--------------------------------------------------------------------------------------
code_list_, code_optimizers, density_grid, density_bitfield = self.load_cache(data)
      
      cache_list=[None, None, None, None, None, None, None, None]  # 初始化为bs个None
      # 循环bs次--------------------------------------------------------------
      for scene_state_single in cache_list:
          if scene_state_single is None:
              # 随机生成以下参数(torch.empty等)
              code_list_.append(self.get_init_code_(None, device))                   # (3,6,128,128)
              density_grid.append(self.get_init_density_grid(None, device))          # (262144)
              density_bitfield.append(self.get_init_density_bitfield(None, device))  #(32768)

# 2.读入图像------------------------------------------------------------------------------------------
concat_cond = None                     # 这里没有使用图像拼接条件
if 'cond_imgs' in data:
    cond_imgs = data['cond_imgs']  #  (8,50,128,128,3) 8个batch,50个视角的图
    cond_intrinsics = data['cond_intrinsics']  # [fx, fy, cx, cy] (bs,50,4)
    cond_poses = data['cond_poses']       # (bs,50,4,4)

# 3. 计算射线(根据pose和intrinc)-------------------------------------------------------------------
 cond_rays_o, cond_rays_d = get_cam_rays(cond_poses, cond_intrinsics, h, w)  # (bs,50,128,128,3)
 dt_gamma_scale = self.train_cfg.get('dt_gamma_scale', 0.0)        # 0.5
 dt_gamma = dt_gamma_scale / cond_intrinsics[..., :2].mean(dim=(-2, -1))    # [0.0038]*(8)

# 4.采样时间步 t------------------------------
t = self.sampler(num_batches).to(device)  # [605, 733, 138, 312, 997, 128, 178, 752]
# 5. 生成噪声
noise = torch.rand().to(device)    # (bs, 18, 128, 128)

# 6.根据时间t,采样噪声
x_t, mean, std = self.q_sample(x_0, t, noise)  # x_0就是上步的code_list_,经过缩放和激活

			    def q_sample(self, x_0, t, noise=None):

			        mean = var_to_tensor(self.sqrt_alphas_bar, t.cpu(), x_0.shape, device)  将输入变量sqrt_alphas转换为张量,并根据索引t提取相应的值    
			        std = var_to_tensor(self.sqrt_one_minus_alphas_bar, t.cpu(), x_0.shape, device)
			        return x_0 * mean + noise * std,  mean,  std

# 7.逆扩散过程-------------------------------------------------------------------------------------------
denoising_output = self.pred_x_0( x_t, t, grad_guide_fn=None, concat_cond=None, cfg=cfg, update_denoising_output=True)
		# 7.1 位置编码------------------------------------------------------------------------
		embedding = self.time_embedding(t)    # bs --> (bs,512) 正余弦编码+conv+silu。具体展开如下:
					embedding = self.blocks(self.sinusodial_embedding(t))
							    def sinusodial_embedding(timesteps, dim, max_period=10000):
							        half = dim // 2        # 128//2
							        freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)   
							        # freqs:64:[1.0, 0.8, 0.7, 0.6...0.0014]
							        args = timesteps[:, None].float() * freqs[None]
							        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
							
							        return embedding       # (bs,128)

		# 7.2 上下采样,用3D卷积和注意力qkv--------------------------------------------------------------
        h, hs = x_t, []                            # (bs,18,128,128)
        if self.concat_cond_channels > 0:
            h = torch.cat([h, concat_cond], dim=1)
        # forward downsample blocks
        for block in self.in_blocks:
            h = block(h, embedding)
            hs.append(h)                           # (bs,512,8,8)

        # forward middle blocks
        h = self.mid_blocks(h, embedding)

        # forward upsample blocks
        for block in self.out_blocks:
            h = block(torch.cat([h, hs.pop()], dim=1), embedding)    #   (bs,128,128,128)
        outputs = self.out(h)                      # (bs,18,128,128)

        return outputs

# 8.  x_0 是输入x_t与预测值denoising_output的加权差:----------------------------------------------------
if self.denoising_mean_mode.upper() == 'V':
    x_0_pred = sqrt_alpha_bar_t * x_t - sqrt_one_minus_alpha_bar_t * denoising_output

# 9.self.ddpm_loss-----------------------------------------------------------------------------------
class DDPMMSELossMod(DDPMLossMod):
    def mse_loss(pred, target):

    return F.mse_loss(pred, target, reduction='none')    # (bs,18,128,128) target就是v_t
    
loss.mean(dim=list(range(1, loss.ndim)=4))    # (bs)
loss_rescaled = loss * weight.to(timesteps.device)[timesteps] * self.weight_scale   # weight_scale=4

quartile = (timesteps / total_timesteps * 4).int()        # [2, 3, 0, 0, 3, 2, 0, 2]
for idx in range(4):
    loss_quartile = reduce_loss(loss[quartile == idx], reduction)   # mean
    log_vars[f'{prefix_name}_quartile_{idx}'] = loss_quartile.item()    # log_vars包含4个loss,求均值得到loss_diffusion
loss_diffusion.backward()

Em seguida, atualize o gradiente em cache:

    if extra_scene_step > 0:
        prior_grad = [code_.grad.data.clone() for code_ in code_list_]        # (8,3,6,128,128)

2.inverse_code: perda de NeRF

# 10.NeRF渲染-----------------------------------------------------------------------------------
for inverse_step_id in range(n_inverse_steps):                  # 循环16if inverse_step_id % self.update_extra_interval == 0:       # 16 的整数倍
        # density_grid, density_bitfield 分别是0初始化的(bs,262144)(bs,32768);经过upgrade后,填充进了sigma
        # update_extra_state 是网格生成过程(同上述无条件生成的3.4节),得到coord与sigma
        self.update_extra_state(decoder, code, density_grid, density_bitfield, iter_density, density_thresh=cfg.get('density_thresh', 0.01))

    inds = raybatch_inds[inverse_step_id % num_raybatch]       # 从 0 到 第200份的射线
    rays_o, rays_d, target_rgbs = self.ray_sample(
        cond_rays_o, cond_rays_d, cond_imgs, n_inverse_rays, sample_inds=inds)   # 选出某份射线:(bs,4096,3)


	## 11.march_rays_train(cuda封装的函数)-------------------------------------------------------------
	xyzs_single, dirs_single, deltas_single, rays_single = march_rays_train(
					                    rays_o_single, rays_d_single, self.bound, density_bitfield_single, 1, grid_size_single, nears_single, fars_single, 
					                    perturb=perturb, align=128, force_all_rays=True, dt_gamma=dt_gamma_single.item(), max_steps=self.max_steps)
	nears, fars = batch_near_far_from_aabb(rays_o, rays_d, self.aabb, self.min_near)
	weights_sum_, depth_, image_ = composite_rays_train(sigmas, rgbs, deltas_, rays_, T_thresh)   # (8,4096)(8,4096)(8,4096,3)
	
	
	## 12.最终的2个损失
	pixel_loss = self.pixel_loss(out_rgbs, target_rgbs, **kwargs) * (scale * 3)     # F.mse_loss, 权重weight=20,求平均   ->7.73
	reg_loss = (code.abs() ** 2).mean()                       # code:特征的大小  -> 3.99e-11
	loss = pixel_loss + reg_loss
	
	## 13. 复制prior_grad的梯度到 code_ (在一个inverse_step中,code_梯度永远来自prior_grad)----------------
    if prior_grad is not None:
        if isinstance(code_, list):
            for code_single_, prior_grad_single in zip(code_, prior_grad):           # 循环8次
                code_single_.grad.copy_(prior_grad_single)      # (3,6,128,128)

    loss.backward()
  1. Atualizar cache
self.save_cache(
    code_list_, code_optimizers,
    density_grid, density_bitfield, data['scene_id'], data['scene_name'])

5.3 Raciocínio de reconstrução esparsa: ssdnerf_cars3v_recons1v

O arquivo de configuração é: ssdnerf_cars3v_recons1v.py Selecione uma música por vez para reconstruir a cena 3D.

with torch.no_grad():
    cond_mode = self.test_cfg.get('cond_mode', 'guide')     # 'guide_optim'
    # 1.val_guide 得到三平面特征code, 以及空间网格特征
    code, density_grid, density_bitfield = self.val_guide(data, **kwargs)    # code:(bs,18,128,128)
            # 1.0 初始化三大变量--------------------------------------------------------------
            density_grid = torch.zeros((num_scenes, self.grid_size ** 3), device=device)        # (8,262144)
            density_bitfield = torch.zeros((num_scenes, self.grid_size ** 3 // 8), dtype=torch.uint8, device=device)    # (8,32768)
            noise = torch.randn((num_scenes, *self.code_size), device)          # (8,3,6,128,128)
            
            # 1.1 扩散75-----------------------------------------------------------------------
		    x_t = noise
		    num_timesteps = self.test_cfg.get('num_timesteps', self.num_timesteps)        # 75
	        langevin_t_range = self.test_cfg.get('langevin_t_range', [0, 1000])           # [0,1000]
		    timesteps = torch.arange(start= 75 - 1, end=-1, step=-(1000/ 75)    # 75[0,1000]的随机数

            for step, t in enumerate(timesteps):                                    # 逆扩散75步,每次更新x_t
                x_t, x_0_pred = self.p_sample_ddim( x_t, t, t_prev, concat_cond=None)

			          # 1.1.1 Unet 通过x_t与t,预测x_0:---------------------------------------------------------
				   	  denoising_output = self.denoising(x_t, t, concat_cond=concat_cond)      # (8,18,128,128)
					  x_0_pred = sqrt_alpha_bar_t * x_t - sqrt_one_minus_alpha_bar_t * denoising_output
					  self.update_extra_state(decoder, x_0_pred, density_grid, density_bitfield,
					                             0, density_thresh=self.test_cfg.get('density_thresh', 0.01))
			          # 1.1.2 利用 x_0_pred 三平面特征,特征空间解码,更新density_grid,density_bitfield-------------------
			          rays_o, rays_d, target_rgbs = self.ray_sample(cond_rays_o, cond_rays_d, 
			                                             cond_imgs, n_inverse_rays, sample_inds=inds)  # (bs,16384,3)
			          # 1.1.3 NeRF 射线解码-----------------------------------------------------------------
			          nears, fars = batch_near_far_from_aabb(rays_o, rays_d, self.aabb, self.min_near=0.2)
			          xyzs_single, dirs_single, deltas_single, rays_single = march_rays_train()
			          sigmas, rgbs, num_points = self.point_decode(xyzs, dirs, code)    # F_grid插值,fc层计算密度/颜色 xyz(bs,93440,3-> (724480)(3495936,3) 8个点数
			          weights_sum, depth, image = batch_composite_rays_train(sigmas, rgbs, deltas, rays, num_points, T_thresh)   # (8,16384)(8,16384)(8,16384,3)
			          # 1.1.4 计算loss--------------------------------------------------------------
			          loss = grad_guide_fn(x_0_pred)
			          pixel_loss = self.pixel_loss(out_rgbs, target_rgbs, **kwargs) * (scale * 3)
			          reg_loss = self.reg_loss(code, **kwargs)
			          loss = pixel_loss + reg_loss
			          grad = torch.autograd.grad(loss, x_t)[0]   # (8,18,128,128) 计算 loss对x_t的梯度
			          torch.set_grad_enabled(False)
			          x_0_pred.detach_()
			          x_0_pred -= grad * ((sqrt_one_minus_alpha_bar_t ** (2 - snr_weight_power * 2))
			              * (sqrt_alpha_bar_t ** (snr_weight_power * 2 - 1))* guidance_gain)   # guidance_gain=52484.8
			          eps_t_pred = (x_t - self.sqrt_alphas_bar[t] * x_0_pred) / self.sqrt_one_minus_alphas_bar[t]
			          pred_sample_direction = np.sqrt(1 - alpha_bar_t_prev - tilde_beta_t * (eta ** 2)) * eps_t_pred    # (8,18,128,128)
			          x_t = np.sqrt(alpha_bar_t_prev) * x_0_pred + pred_sample_direction                             # (8,18,128,128)

    # 2. 继续逆扩散,更新特征---------------------------------------------------------------------------
    code, density_grid, density_bitfield = self.val_optim(data, code_=self.code_activation.inverse(code).requires_grad_(True),
                                                          density_grid, density_bitfield=density_bitfield, **kwargs)
            # 2.1 根据pose,计算射线-----------------------------------------------------------------------
	        cond_imgs = data['cond_imgs']  # (num_scenes, num_imgs, h, w, 3)
	        cond_intrinsics = data['cond_intrinsics']  # (num_scenes, num_imgs, 4), in [fx, fy, cx, cy]
	        cond_poses = data['cond_poses']
	
	        cond_rays_o, cond_rays_d = get_cam_rays(cond_poses, cond_intrinsics, h, w)     # (bs,1,128,128,3)
         
            for inverse_step_id in range(n_inverse_steps):                             # 25个inverse步骤
                loss, log_vars = diffusion(self.code_diff_pr(code), return_loss=True, concat_cond=None)
                # 上面是diff_train,采样t,Unet去噪, 计算4个ddpm_mes损失
                loss.backward()
                
                prior_grad = code_.grad.data.clone()
                self.inverse_code(decoder, cond_imgs, cond_rays_o, cond_rays_d, dt_gamma=dt_gamma, cfg code_=code_, density_grid, density_bitfield, prior_grad=prior_grad)

                     # inverse_code包含以下步骤:
                     for inverse_step_id in range(n_inverse_steps):        # 循环4if inverse_step_id % self.update_extra_interval == 0:
                            self.update_extra_state(decoder, code, density_grid, density_bitfield,iter_density, 
                                                    density_thresh=cfg.get('density_thresh', 0.01))    # 生成网格,F特征插值、base_net更新sigma

	                     rays_o, rays_d, target_rgbs = self.ray_sample( cond_rays_o, cond_rays_d, cond_imgs, n_inverse_rays, sample_inds=inds)        # (bs, 16384, 3)
	                     out_rgbs, loss, loss_dict = self.loss( decoder, code, density_bitfield, target_rgbs, rays_o, rays_d)
	                     code_.grad.copy_(prior_grad)
	                     loss.backward()
 # 3.-------------------------------------------------------------------------------------------
 log_vars, pred_imgs = self.eval_and_viz(data, decoder, code, density_bitfield, viz_dir=viz_dir, cfg)
        image, depth = self.render( decoder, code, density_bitfield, h, w, test_intrinsics, test_poses, cfg=cfg)      #code:(8,3,6,128,128)bitfield:(8,32768) pose:(8,251,4,4)

A seguir está uma expansão da Parte 3: self.render

# 输入数据的num_imgs有250张
test_intrinsics = data['test_intrinsics']  # (num_scenes, num_imgs, 4), in [fx, fy, cx, cy]
test_poses = data['test_poses']
test_imgs = data['test_imgs']  # (num_scenes, num_imgs, h, w, 3)

outputs = decode(batch_near_far_from_aabb、march_rays、self.point_decode等,循环self.max_steps=256)
        results = dict(
            weights_sum=weights_sum,          # bs个(4096000)
            depth=depth,                      # bs个(4096000)
            image=image)                      # bs个(40960003)

weights = torch.stack(outputs['weights_sum'], dim=0) if num_scenes > 1 else outputs['weights_sum'][0]    # 01
rgbs = (torch.stack(outputs['image'], dim=0) if num_scenes > 1 else outputs['image'][0]) \
      + self.bg_color * (1 - weights.unsqueeze(-1))
depth = torch.stack(outputs['depth'], dim=0) if num_scenes > 1 else outputs['depth'][0]
pred_imgs = image.permute(0, 1, 4, 2, 3).reshape(8 * 250, 3, h, w).clamp(min=0, max=1)

# 测试生成的全视角图像,和数据集真实图像的指标
test_psnr = eval_psnr(pred_imgs, target_imgs)   # bs*250个数
test_ssim = eval_ssim_skimage(pred_imgs, target_imgs, data_range=1)

# 最后是保存图像(与真实图像拼接,便于对比)在work_dir文件夹中,此处忽略。。。。。。。。。。。。。。

5.4 Treinamento de reconstrução esparsa: ssdnerf_cars3v_recons1v

O arquivo de configuração é: ssdnerf_cars3v_recons1v.py Cada vez, 3 perspectivas são selecionadas para treinamento de reconstrução.

# 1.载入缓存.若无缓存,初始化(生成)code等三大变量:------------------------------------------------
code_list_, code_optimizers, density_grid, density_bitfield = self.load_cache(data)
code = self.code_activation(torch.stack(code_list_, dim=0), update_stats=True)

# 2.载入图像----------------------------------------------------------------------------
cond_imgs = data['cond_imgs']  # (num_scenes, num_imgs, h, w, 3)     (8,3,128,128,3)
cond_intrinsics = data['cond_intrinsics']  # (num_scenes, num_imgs, 4), in [fx, fy, cx, cy] (8,3,4)
cond_poses = data['cond_poses']       # (8,3,4,4)

num_scenes, num_imgs, h, w, _ = cond_imgs.size()
# (num_scenes, num_imgs, h, w, 3)
cond_rays_o, cond_rays_d = get_cam_rays(cond_poses, cond_intrinsics, h, w)  # (8,3,128,128,3)
dt_gamma_scale = self.train_cfg.get('dt_gamma_scale', 0.0)        # 0.5
# (num_scenes,)
dt_gamma = dt_gamma_scale / cond_intrinsics[..., :2].mean(dim=(-2, -1))    # [0.0038]*(8)

# 3. 扩散模型损失
loss_diffusion, log_vars = diffusion(self.code_diff_pr(code), concat_cond=None, return_loss=True,
                                     x_t_detach=x_t_detach, cfg=self.train_cfg)
                t = self.sampler(num_batches).to(device)              # [630, 757, 989, 452,...] bs个
                noise = torch.randn((num_batches, *image_shape))
                x_t, mean, std = self.q_sample(x_0, t, noise)         # 从x_0 采样到x_t 
                denoising_output = self.denoising(x_t, t, concat_cond=concat_cond)      # Unet逆扩散过程(8,18,128,128)

loss_diffusion.backward()                

# 4. 更新 先验梯度-------------------------------------------------------------------------------------
prior_grad = [code_.grad.data.clone() for code_ in code_list_]        # (8,3,6,128,128)

# 5.self.inverse_code:NeRF渲染:----------------------------------------------------------------------
raybatch_inds, num_raybatch = self.get_raybatch_inds(cond_imgs, n_inverse_rays)  # (8,3,128,12834096 -> 随机分成 12份的 (8,4096)
self.update_extra_state(decoder, code, density_grid, density_bitfield,
                        iter_density, density_thresh=cfg.get('density_thresh', 0.01))

nears, fars = batch_near_far_from_aabb(rays_o, rays_d, self.aabb, self.min_near)
xyzs_single, dirs_single, deltas_single, rays_single = march_rays_train( rays_o_single, rays_d_single, 
                                 self.bound, density_bitfield_single, 1, grid_size_single, nears_single, fars_single,  perturb=True)


sigmas, rgbs, num_points = self.point_decode(xyzs, dirs, code)    # F_grid插值,fc层计算密度/颜色 (8,438144,3-> (3495424)(3495424,3) 8个点数
            weights_sum, depth, image = batch_composite_rays_train(sigmas, rgbs, deltas, rays, num_points, T_thresh)   # (8,4096)(8,4096)(8,4096,3)

pixel_loss = self.pixel_loss(out_rgbs, target_rgbs, **kwargs) * (scale * 3)     # F.mse_loss, 权重weight=20,求平均   ->7.73
reg_loss = self.reg_loss(code, **kwargs)
loss = pixel_loss + reg_loss

for code_single_, prior_grad_single in zip(code_, prior_grad):           # 循环8次
    code_single_.grad.copy_(prior_grad_single)      # (3,6,128,128)
loss.backward()

# 6.更新 density_grid, density_bitfield------------------------------------------------------------
self.update_extra_state( decoder, code, density_grid, density_bitfield,
                0, density_thresh=self.train_cfg.get('density_thresh', 0.01))


# 7.loss_decoder = self.loss_decoder(decoder, code, density_bitfield, cond_rays_o, cond_rays_d,
                                     cond_imgs, dt_gamma, cfg=self.train_cfg) #-----------------------
    # decoder_loss作为第三个损失,与第二部损失相同。内容如下:
	nears, fars = batch_near_far_from_aabb(rays_o, rays_d, self.aabb, self.min_near)
	for i in range(batchsize):
		xyzs_single, dirs_single, deltas_single, rays_single = march_rays_train(
							                   rays_o_single, rays_d_single, self.bound, density_bitfield_single,
							                   1, grid_size_single, nears_single, fars_single,
							                   perturb=perturb, align=128, force_all_rays=True,
							                   dt_gamma_single.item(), max_steps=256)      # (4096,3->(435968,3)(435968,3)(435968,2)
	    sigmas, rgbs, num_points = self.point_decode(xyzs, dirs, code)    # F_grid插值,fc层计算密度/颜色 (8,447212,3-> (3495936)(3495936,3) 8个点数
	    weights_sum, depth, image = batch_composite_rays_train(sigmas, rgbs, deltas, rays, num_points, T_thresh)   # (8,4096)(8,4096)(8,4096,3)

# 8.更新 decoder 梯度------------------------------------------------------------------------------
for code_, prior_grad_single in zip(code_list_, prior_grad):
    code_.grad.copy_(prior_grad_single)
loss_decoder.backward()

# 9.缓存cache-------------------------------------------------------------------------------------
self.save_cache( code_list_, code_optimizers, density_grid, density_bitfield, data['scene_id'], data['scene_name'])

# 10.验证-----------------------------------------------------------------------------------------
train_psnr = eval_psnr(out_rgbs, target_rgbs)

* Funções encapsuladas .cuda

lib/ops/raymarching/scr/raymarching.h contém as seguintes funções (definidas especificamente no arquivo raymarching.cu)

#pragma once

#include <stdint.h>
#include <torch/torch.h>

void near_far_from_aabb(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor aabb, const uint32_t N, const float min_near, at::Tensor nears, at::Tensor fars);
void sph_from_ray(const at::Tensor rays_o, const at::Tensor rays_d, const float radius, const uint32_t N, at::Tensor coords);
void morton3D(const at::Tensor coords, const uint32_t N, at::Tensor indices);
void morton3D_invert(const at::Tensor indices, const uint32_t N, at::Tensor coords);
void packbits(const at::Tensor grid, const uint32_t N, const float density_thresh, at::Tensor bitfield);

void march_rays_train(const at::Tensor rays_o, const at::Tensor rays_d, const at::Tensor grid, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t N, const uint32_t C, const uint32_t H, const uint32_t M, const at::Tensor nears, const at::Tensor fars, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor rays, at::Tensor counter, at::Tensor noises);
void composite_rays_train_forward(const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor weights_sum, at::Tensor depth, at::Tensor image);
void composite_rays_train_backward(const at::Tensor grad_weights_sum, const at::Tensor grad_image, const at::Tensor sigmas, const at::Tensor rgbs, const at::Tensor deltas, const at::Tensor rays, const at::Tensor weights_sum, const at::Tensor image, const uint32_t M, const uint32_t N, const float T_thresh, at::Tensor grad_sigmas, at::Tensor grad_rgbs);

void march_rays(const uint32_t n_alive, const uint32_t n_step, const at::Tensor rays_alive, const at::Tensor rays_t, const at::Tensor rays_o, const at::Tensor rays_d, const float bound, const float dt_gamma, const uint32_t max_steps, const uint32_t C, const uint32_t H, const at::Tensor grid, const at::Tensor nears, const at::Tensor fars, at::Tensor xyzs, at::Tensor dirs, at::Tensor deltas, at::Tensor noises);
void composite_rays(const uint32_t n_alive, const uint32_t n_step, const float T_thresh, at::Tensor rays_alive, at::Tensor rays_t, at::Tensor sigmas, at::Tensor rgbs, at::Tensor deltas, at::Tensor weights_sum, at::Tensor depth, at::Tensor image);

limitação

     Atualmente, este método depende dos parâmetros da câmera Ground Truth durante treinamento e teste. Trabalhos futuros podem explorar modelos invariantes à transformação. Além disso, à medida que o tempo de treinamento aumenta, a difusão anterior torna-se descontínua, afetando a generalização . Embora a parada antecipada seja usada temporariamente, um melhor projeto de rede ou um conjunto maior de dados de treinamento pode resolver fundamentalmente esse problema.


expandir

1. Indicadores FID, KID, LPIPS

Frechet Inception Distance (FID) é uma métrica utilizada para avaliar a qualidade das imagens geradas por GANs. Baseia-se na distância de Fréchet entre a imagem gerada e a imagem real , que mede a semelhança entre as duas distribuições de imagens. Quanto menor o FID, melhor será a qualidade da imagem gerada.

Insira a descrição da imagem aqui

LPIPS : Extraído do artigo "A eficácia irracional de recursos profundos como uma métrica perceptual".Quanto
Insira a descrição da imagem aqui
menor o valor, mais semelhantes são as duas imagens. As duas entradas são enviadas para a rede neural F (pode ser VGG, Alexnet, Squeezenet) para extração de características, a saída de cada camada é ativada e normalizada, e então a distância L2 é calculada após a multiplicação dos pesos da camada w.
Insira a descrição da imagem aqui

Kernel Inception Distance (KID) é outra métrica para avaliar a qualidade das imagens geradas por GANs. É baseado na distância da função kernel entre a imagem gerada e a imagem real , que mede a diferença entre a imagem real e a imagem gerada mapeando os vetores de recursos em um espaço de alta dimensão e calculando a distância de Frechet da matriz do kernel entre eles. KID também é uma medida da qualidade da imagem gerada, semelhante ao FID.

Para implementação de código específico, consulte https://blog.51cto.com/u_16175458/6906283

2. Regularização da TV

Regularização de TV, o nome completo é Regularização de Variação Total. A regularização da TV consegue suavizar a imagem minimizando a magnitude do gradiente da imagem.

Especificamente, para uma imagem bidimensional, a regularização de TV pode obter suavização minimizando a magnitude do gradiente da imagem . Isso suprime ruídos e detalhes na imagem. Torna a imagem mais suave. A regularização de TV é geralmente aplicada ao termo de regularização de problemas de otimização para equilibrar a relação entre ajuste de dados e suavidade.

Para métodos como NeRF (Neural Radiance Fields), que são usados ​​principalmente para reconstrução tridimensional, a regularização TV ajuda a melhorar a qualidade dos resultados da reconstrução. Isso ocorre porque na reconstrução 3D, devido a fatores como dispersão de dados e ruído, os resultados da reconstrução geralmente contêm detalhes e ruído desnecessários .

Além disso, a regularização da TV também pode ajudar a fortalecer as restrições ao modelo de aprendizagem profunda e ajudar a melhorar a capacidade de generalização e a capacidade anti-ruído do modelo. Portanto, para tarefas de reconstrução 3D como NeRF, a aplicação da regularização TV ajuda a melhorar a qualidade dos resultados da reconstrução e a aumentar a robustez do modelo.

3. Renderização NeRF (dos parâmetros internos e externos da câmera aos raios)

raios_o, raios_d = get_cam_rays(poses, intrínsecos, h, w)
analisa principalmente a função get_cam_rays, localizada em lib/core/utils/nerf_utils.py

directions = get_ray_directions(h, w, intrinsics, norm=False)  # (8, 251, h, w,3)

      def get_ray_directions():
	    batch_size = intrinsics.shape[:-1]
	    x = torch.linspace(0.5, w - 0.5, w, device=device)
	    y = torch.linspace(0.5, h - 0.5, h, device=device)
	    
	    # 相对坐标,除以焦距,就是射线方向----------------------------------------------
	    directions_xy = torch.stack(
	        [((x - intrinsics[..., 2:3]) / intrinsics[..., 0:1])[..., None, :].expand(*batch_size, h, w),
	         ((y - intrinsics[..., 3:4]) / intrinsics[..., 1:2])[..., :, None].expand(*batch_size, h, w)], dim=-1)
	    # (8,251,128,128, 2)
	    directions = F.pad(directions_xy, [0, 1], mode='constant', value=1.0)
	    
	    return directions


rays_o, rays_d = get_rays(directions, c2w, norm=True)

       def get_rays(directions, c2w, norm=False):
		   rays_d = directions @ c2w[..., None, :3, :3].transpose(-1, -2)   # (8,251,128,128,3)
		   rays_o = c2w[..., None, None, :3, 3].expand(rays_d.shape)        # (8,251,128,128,3)
		   if norm:
		       rays_d = F.normalize(rays_d, dim=-1)                         # 坐标归一化到 -11之间
		   return rays_o, rays_d

Acho que você gosta

Origin blog.csdn.net/qq_45752541/article/details/134984758
Recomendado
Clasificación