4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00

Fix SSIM memory leak

This commit is contained in:
Matt Tancik 2022-07-05 11:58:32 -07:00 committed by Matthew Tancik
parent fefbad2592
commit ced02a00be
5 changed files with 15 additions and 10 deletions

View File

@ -20,8 +20,9 @@ from typing import Dict, List
import torch
from torch.nn import Parameter
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.functional import structural_similarity_index_measure
from pyrad.fields.modules.field_heads import FieldHeadNames
from pyrad.fields.instant_ngp_field import field_implementation_to_class
@ -81,7 +82,7 @@ class NGPGraph(Graph):
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self) -> Dict[str, List[Parameter]]:

View File

@ -21,8 +21,9 @@ from typing import Dict, List
import torch
from torch.nn import Parameter
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.functional import structural_similarity_index_measure
from pyrad.fields.modules.encoding import NeRFEncoding
from pyrad.fields.modules.field_heads import FieldHeadNames
@ -85,7 +86,7 @@ class MipNerfGraph(Graph):
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self) -> Dict[str, List[Parameter]]:

View File

@ -21,8 +21,9 @@ from typing import Dict, List
import torch
from torch.nn import Parameter
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.functional import structural_similarity_index_measure
from pyrad.fields.modules.encoding import NeRFEncoding
from pyrad.fields.modules.field_heads import FieldHeadNames
@ -89,7 +90,7 @@ class MipNerf360Graph(Graph):
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self) -> Dict[str, List[Parameter]]:

View File

@ -17,8 +17,9 @@ NeRF-W (NeRF in the wild) implementation.
"""
import torch
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.functional import structural_similarity_index_measure
from pyrad.fields.modules.encoding import NeRFEncoding
from pyrad.fields.modules.field_heads import FieldHeadNames
@ -92,7 +93,7 @@ class NerfWGraph(Graph):
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self):

View File

@ -21,8 +21,9 @@ from typing import Dict, List
import torch
from torch.nn import Parameter
from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics import PeakSignalNoiseRatio
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.functional import structural_similarity_index_measure
from pyrad.fields.modules.encoding import NeRFEncoding
from pyrad.fields.modules.field_heads import FieldHeadNames
@ -94,7 +95,7 @@ class NeRFGraph(Graph):
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
def get_param_groups(self) -> Dict[str, List[Parameter]]: