Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне
assets/conveyors (274 МБ) - ленты и угловая секция NVIDIA, на которые ссылается сцена относительным путём. Раньше исключались как перекачиваемые, но без них сцена не композится из коробки. cv/ - код стереодвижков, которые вызывает control_test, без весов: * defom-stereo - рабочий бейзлайн (DEFOM vitl, вход 480, iters 24) * crestereo - второй движок, точнее по габаритам (MAE 23.5 против 32.8 мм) * fast-foundationstereo - проверялся, в бейзлайн не вошёл * circular_section.py - показатель кругового сечения, перенесён в measure_plane.py: выравнивает облако по СОБСТВЕННЫМ главным осям и режет на пяти высотах вдоль каждой. Три самодельные версии (мировые оси, одно сечение) давали хуже; результаты проверки на эталонной геометрии - в circular_section_results.json Веса по-прежнему не в репозитории - источники в MODELS.md. Наборы кадров прежних прогонов (cv/flow_*, 1.26 ГБ) исключены: это выход, а не исходники. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
dasha_f
+445
@@ -0,0 +1,445 @@
|
||||
import torch,pdb,logging,timm
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import sys,os
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../')
|
||||
from core.update import BasicSelectiveMultiUpdateBlock
|
||||
from core.extractor import ContextNetSharedBackbone, Feature
|
||||
from core.geometry import Combined_Geo_Encoding_Volume
|
||||
from core.submodule import (
|
||||
BasicConv, Conv3dNormActReduced, ResnetBasicBlock3D, BasicConv_IN, Conv2x,
|
||||
FeatureAtt, CostVolumeDisparityAttention, SpatialAttentionExtractor,
|
||||
ChannelAttentionEnhancement, disparity_regression, context_upsample,
|
||||
build_gwc_volume_optimized_pytorch1, build_gwc_volume_triton,
|
||||
build_concat_volume_optimized_pytorch1, build_concat_volume_optimized_pytorch,
|
||||
)
|
||||
from core.utils.utils import InputPadder
|
||||
import Utils as U
|
||||
import time
|
||||
|
||||
sys.modules['foundation_stereo_ori'] = sys.modules['core']
|
||||
sys.modules['foundation_stereo_ori.submodule'] = sys.modules['core.submodule']
|
||||
sys.modules['foundation_stereo_ori.extractor'] = sys.modules['core.extractor']
|
||||
sys.modules['foundation_stereo_ori.update'] = sys.modules['core.update']
|
||||
sys.modules['foundation_stereo_ori.foundation_stereo'] = sys.modules['core.foundation_stereo']
|
||||
class FoundationStereo(nn.Module):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def normalize_image(img):
|
||||
'''
|
||||
@img: (B,C,H,W) in range 0-255, RGB order
|
||||
'''
|
||||
mean = img.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
|
||||
std = img.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
|
||||
return (img/255.0 - mean) / std
|
||||
|
||||
|
||||
class hourglass(nn.Module):
|
||||
def __init__(self, cfg, in_channels, feat_dims=None):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
self.conv1 = nn.Sequential(BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3,
|
||||
padding=1, stride=2, dilation=1),
|
||||
Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
|
||||
|
||||
self.conv2 = nn.Sequential(BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3,
|
||||
padding=1, stride=2, dilation=1),
|
||||
Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17))
|
||||
|
||||
self.conv3 = nn.Sequential(BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3,
|
||||
padding=1, stride=2, dilation=1),
|
||||
Conv3dNormActReduced(in_channels*6, in_channels*6, kernel_size=3, kernel_disp=17))
|
||||
|
||||
|
||||
self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True,
|
||||
relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
|
||||
|
||||
self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True,
|
||||
relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
|
||||
|
||||
self.conv1_up = BasicConv(in_channels*2, in_channels, deconv=True, is_3d=True, bn=True,
|
||||
relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
|
||||
self.conv_out = nn.Sequential(
|
||||
Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
|
||||
Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
|
||||
)
|
||||
|
||||
self.agg_0 = nn.Sequential(BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1),
|
||||
Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),
|
||||
Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),)
|
||||
|
||||
self.agg_1 = nn.Sequential(BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1),
|
||||
Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17),
|
||||
Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
|
||||
|
||||
self.atts = nn.ModuleDict({
|
||||
"4": CostVolumeDisparityAttention(d_model=in_channels, nhead=4, dim_feedforward=in_channels, norm_first=False, num_transformer=4, max_len=self.cfg['max_disp']//16),
|
||||
})
|
||||
self.conv_patch = nn.Sequential(
|
||||
nn.Conv3d(in_channels, in_channels, kernel_size=4, stride=4, padding=0, groups=in_channels),
|
||||
nn.BatchNorm3d(in_channels),
|
||||
)
|
||||
self.feature_att_8 = FeatureAtt(in_channels*2, feat_dims[1])
|
||||
self.feature_att_16 = FeatureAtt(in_channels*4, feat_dims[2])
|
||||
self.feature_att_32 = FeatureAtt(in_channels*6, feat_dims[3])
|
||||
self.feature_att_up_16 = FeatureAtt(in_channels*4, feat_dims[2])
|
||||
self.feature_att_up_8 = FeatureAtt(in_channels*2, feat_dims[1])
|
||||
|
||||
self.post32_to_16 = None
|
||||
self.post16_to_8 = None
|
||||
self.post8_to_4 = None
|
||||
|
||||
def forward(self, x, features):
|
||||
conv1 = self.conv1(x)
|
||||
conv1 = self.feature_att_8(conv1, features[1])
|
||||
|
||||
conv2 = self.conv2(conv1)
|
||||
conv2 = self.feature_att_16(conv2, features[2])
|
||||
|
||||
conv3 = self.conv3(conv2)
|
||||
conv3 = self.feature_att_32(conv3, features[3])
|
||||
if not hasattr(self, 'post32_to_16') or self.post32_to_16 is None:
|
||||
conv3_up = self.conv3_up(conv3)
|
||||
conv2 = torch.cat((conv3_up, conv2), dim=1)
|
||||
conv2 = self.agg_0(conv2)
|
||||
conv2 = self.feature_att_up_16(conv2, features[2])
|
||||
else:
|
||||
conv2 = self.post32_to_16(conv2, conv3, features[2])
|
||||
|
||||
if not hasattr(self, 'post16_to_8') or self.post16_to_8 is None:
|
||||
conv2_up = self.conv2_up(conv2)
|
||||
conv1 = torch.cat((conv2_up, conv1), dim=1)
|
||||
conv1 = self.agg_1(conv1)
|
||||
conv1 = self.feature_att_up_8(conv1, features[1])
|
||||
else:
|
||||
conv1 = self.post16_to_8(conv1, conv2, features[1])
|
||||
|
||||
conv = self.conv1_up(conv1)
|
||||
if not hasattr(self, 'post8_to_4') or self.post8_to_4 is None:
|
||||
x = self.conv_patch(x)
|
||||
x = self.atts["4"](x)
|
||||
x = F.interpolate(x, scale_factor=4, mode='trilinear', align_corners=False)
|
||||
conv = conv + x
|
||||
conv = self.conv_out(conv)
|
||||
else:
|
||||
conv = self.post8_to_4(x, conv)
|
||||
|
||||
return conv
|
||||
|
||||
|
||||
class FastFoundationStereo(nn.Module):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
self.dtype = torch.float32
|
||||
|
||||
context_dims = args.hidden_dims
|
||||
self.cv_group = args.get('cv_group', 8)
|
||||
self.concat_channel = 24
|
||||
volume_dim = args.get('volume_dim', 28)
|
||||
self.volume_dim = volume_dim
|
||||
self.update_block = BasicSelectiveMultiUpdateBlock(self.args, self.args.hidden_dims[0], volume_dim=volume_dim)
|
||||
self.sam = SpatialAttentionExtractor()
|
||||
self.cam = ChannelAttentionEnhancement(self.args.hidden_dims[0])
|
||||
self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, kernel_size=3, padding=3//2) for i in range(self.args.n_gru_layers)])
|
||||
self.feature = Feature(args)
|
||||
self.proj_cmb = nn.Conv2d(self.feature.d_out[0], self.concat_channel//2, kernel_size=1, padding=0)
|
||||
self.cnet = ContextNetSharedBackbone(args, c04=self.feature.d_out[0], c08=self.feature.d_out[1], c16=self.feature.d_out[2], output_dim=[args.hidden_dims, context_dims])
|
||||
|
||||
self.stem_2 = nn.Sequential(
|
||||
BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1),
|
||||
nn.Conv2d(32, 32, 3, 1, 1, bias=False),
|
||||
nn.InstanceNorm2d(32), nn.ReLU()
|
||||
)
|
||||
self.spx_2_gru = Conv2x(32, 32, deconv=True, bn=False, concat=True)
|
||||
self.spx_gru = nn.Sequential(
|
||||
nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1),
|
||||
)
|
||||
|
||||
self.corr_stem = nn.Sequential(
|
||||
nn.Conv3d(self.proj_cmb.out_channels*2+self.cv_group, volume_dim, kernel_size=1),
|
||||
BasicConv(volume_dim, volume_dim, kernel_size=3, padding=1, is_3d=True),
|
||||
ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
|
||||
ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
|
||||
)
|
||||
self.corr_feature_att = FeatureAtt(volume_dim, self.feature.d_out[0])
|
||||
self.cost_agg = hourglass(cfg=self.args, in_channels=volume_dim, feat_dims=self.feature.d_out)
|
||||
self.classifier = nn.Sequential(
|
||||
BasicConv(volume_dim, volume_dim//2, kernel_size=3, padding=1, is_3d=True),
|
||||
ResnetBasicBlock3D(volume_dim//2, volume_dim//2, kernel_size=3, stride=1, padding=1),
|
||||
nn.Conv3d(volume_dim//2, 1, kernel_size=7, padding=3),
|
||||
)
|
||||
|
||||
r = self.args.corr_radius
|
||||
dx = torch.arange(-r, r+1, requires_grad=False, dtype=torch.int8).reshape(1, 1, 2*r+1, 1)
|
||||
self.register_buffer("dx", dx)
|
||||
|
||||
|
||||
def upsample_disp(self, disp, mask_feat_4, stem_2x):
|
||||
with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
|
||||
xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 1/2 resolution
|
||||
spx_pred = self.spx_gru(xspx)
|
||||
spx_pred = F.softmax(spx_pred, 1)
|
||||
up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1)
|
||||
return up_disp.to(self.dtype)
|
||||
|
||||
|
||||
def forward(self, image1, image2, iters=12, test_mode=False, low_memory=False, init_disp=None, profile=False, optimize_build_volume='pytorch1'):
|
||||
""" Estimate disparity between pair of frames """
|
||||
B,C,H,W = image1.shape
|
||||
low_memory = low_memory or (self.args.get('low_memory', False))
|
||||
image1 = normalize_image(image1)
|
||||
image2 = normalize_image(image2)
|
||||
with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
|
||||
out = self.feature(torch.cat([image1, image2], dim=0))
|
||||
features_left = [o[:B] for o in out]
|
||||
features_right = [o[B:] for o in out]
|
||||
stem_2x = self.stem_2(image1)
|
||||
|
||||
if optimize_build_volume=='pytorch1':
|
||||
gwc_volume = build_gwc_volume_optimized_pytorch1(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
|
||||
elif optimize_build_volume=='triton':
|
||||
gwc_volume = build_gwc_volume_triton(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
|
||||
else:
|
||||
raise RuntimeError(f"Invalid optimize_build_volume: {optimize_build_volume}")
|
||||
|
||||
left_tmp = self.proj_cmb(features_left[0])
|
||||
right_tmp = self.proj_cmb(features_right[0])
|
||||
concat_volume = build_concat_volume_optimized_pytorch1(left_tmp, right_tmp, maxdisp=self.args.max_disp//4)
|
||||
del left_tmp, right_tmp
|
||||
comb_volume = torch.cat([gwc_volume, concat_volume], dim=1)
|
||||
del concat_volume, gwc_volume
|
||||
|
||||
comb_volume = self.corr_stem(comb_volume)
|
||||
comb_volume = self.corr_feature_att(comb_volume, features_left[0])
|
||||
comb_volume = self.cost_agg(comb_volume, features_left)
|
||||
|
||||
# Init disp from geometry encoding volume
|
||||
logits = self.classifier(comb_volume).squeeze(1)
|
||||
prob = F.softmax(logits, dim=1)
|
||||
if init_disp is None:
|
||||
init_disp = disparity_regression(prob, self.args.max_disp//4)
|
||||
|
||||
cnet_list = self.cnet(features_left[0], features_left[1], features_left[2])
|
||||
cnet_list = list(cnet_list)
|
||||
net_list = [torch.tanh(x[0]) for x in cnet_list]
|
||||
inp_list = [torch.relu(x[1]) for x in cnet_list]
|
||||
inp_list = [self.cam(x) * x for x in inp_list]
|
||||
att = [self.sam(x) for x in inp_list]
|
||||
|
||||
geo_fn = Combined_Geo_Encoding_Volume(features_left[0].to(self.dtype), features_right[0].to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels)
|
||||
b, c, h, w = features_left[0].shape
|
||||
coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1)
|
||||
disp = init_disp.to(self.dtype)
|
||||
disp_preds = []
|
||||
|
||||
del comb_volume, features_left, features_right, cnet_list
|
||||
|
||||
# GRUs iterations to update disparity (1/4 resolution)
|
||||
for itr in range(iters):
|
||||
disp = disp.detach()
|
||||
geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=low_memory)
|
||||
with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
|
||||
net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att)
|
||||
|
||||
disp = disp + delta_disp.to(self.dtype)
|
||||
if test_mode and itr < iters-1:
|
||||
continue
|
||||
|
||||
# upsample predictions
|
||||
disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype))
|
||||
disp_preds.append(disp_up)
|
||||
|
||||
|
||||
if test_mode:
|
||||
return disp_up
|
||||
|
||||
return init_disp, disp_preds
|
||||
|
||||
|
||||
def run_hierachical(self, image1, image2, iters=12, test_mode=False, low_memory=False, small_ratio=0.5):
|
||||
B,_,H,W = image1.shape
|
||||
img1_small = F.interpolate(image1, scale_factor=small_ratio, align_corners=False, mode='bilinear')
|
||||
img2_small = F.interpolate(image2, scale_factor=small_ratio, align_corners=False, mode='bilinear')
|
||||
padder = InputPadder(img1_small.shape[-2:], divis_by=32, force_square=False)
|
||||
img1_small, img2_small = padder.pad(img1_small, img2_small)
|
||||
disp_small = self.forward(img1_small, img2_small, test_mode=True, iters=iters, low_memory=low_memory)
|
||||
disp_small = padder.unpad(disp_small)
|
||||
disp_small_up = F.interpolate(disp_small, size=(H,W), mode='bilinear', align_corners=True) * 1/small_ratio
|
||||
disp_small_up = disp_small_up.clip(0, None)
|
||||
|
||||
padder = InputPadder(image1.shape[-2:], divis_by=32, force_square=False)
|
||||
image1, image2, disp_small_up = padder.pad(image1, image2, disp_small_up)
|
||||
disp_small_up += padder._pad[0]
|
||||
init_disp = F.interpolate(disp_small_up, scale_factor=0.25, mode='bilinear', align_corners=True) * 0.25 # Init disp will be 1/4
|
||||
disp = self.forward(image1, image2, iters=iters, test_mode=test_mode, low_memory=low_memory, init_disp=init_disp)
|
||||
disp = padder.unpad(disp)
|
||||
return disp
|
||||
|
||||
FoundationStereoLite = FastFoundationStereo
|
||||
|
||||
|
||||
class TrtFeatureRunner(nn.Module):
|
||||
def __init__(self, model):
|
||||
super().__init__()
|
||||
self.feature = model.feature
|
||||
self.stem_2 = model.stem_2
|
||||
|
||||
def forward(self, image1, image2):
|
||||
image1 = normalize_image(image1)
|
||||
image2 = normalize_image(image2)
|
||||
B = len(image1)
|
||||
out = self.feature(torch.cat([image1, image2], dim=0))
|
||||
features_left = [o[:B] for o in out]
|
||||
features_right = [o[B:] for o in out]
|
||||
stem_2x = self.stem_2(image1)
|
||||
return *features_left, features_right[0], stem_2x
|
||||
|
||||
|
||||
class TrtPostRunner(nn.Module):
|
||||
def __init__(self, model):
|
||||
super().__init__()
|
||||
self.args = model.args
|
||||
self.dtype = model.dtype
|
||||
self.register_buffer("dx", model.dx)
|
||||
self.proj_cmb = model.proj_cmb
|
||||
self.corr_stem = model.corr_stem
|
||||
self.corr_feature_att = model.corr_feature_att
|
||||
self.cost_agg = model.cost_agg
|
||||
self.classifier = model.classifier
|
||||
self.update_block = model.update_block
|
||||
self.sam = model.sam
|
||||
self.cam = model.cam
|
||||
self.feature = model.feature
|
||||
self.spx_2_gru = model.spx_2_gru
|
||||
self.spx_gru = model.spx_gru
|
||||
self.cnet = model.cnet
|
||||
|
||||
def upsample_disp(self, disp, mask_feat_4, stem_2x):
|
||||
with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
|
||||
xspx = self.spx_2_gru(mask_feat_4, stem_2x) # 1/2 resolution
|
||||
spx_pred = self.spx_gru(xspx)
|
||||
spx_pred = F.softmax(spx_pred, 1)
|
||||
up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1)
|
||||
return up_disp.to(self.dtype)
|
||||
|
||||
|
||||
def forward(self, features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume):
|
||||
features_left = [features_left_04, features_left_08, features_left_16, features_left_32]
|
||||
left_tmp = self.proj_cmb(features_left_04)
|
||||
right_tmp = self.proj_cmb(features_right_04)
|
||||
concat_volume = build_concat_volume_optimized_pytorch(left_tmp, right_tmp, maxdisp=self.args.max_disp//4)
|
||||
del left_tmp, right_tmp
|
||||
comb_volume = torch.cat([gwc_volume, concat_volume], dim=1)
|
||||
del concat_volume, gwc_volume
|
||||
comb_volume = self.corr_stem(comb_volume)
|
||||
comb_volume = self.corr_feature_att(comb_volume, features_left_04)
|
||||
comb_volume = self.cost_agg(comb_volume, features_left)
|
||||
|
||||
# Init disp from geometry encoding volume
|
||||
logits = self.classifier(comb_volume).squeeze(1)
|
||||
prob = F.softmax(logits, dim=1)
|
||||
init_disp = disparity_regression(prob, self.args.max_disp//4)
|
||||
|
||||
cnet_list = self.cnet(features_left[0], features_left[1], features_left[2])
|
||||
cnet_list = list(cnet_list)
|
||||
net_list = [torch.tanh(x[0]) for x in cnet_list]
|
||||
inp_list = [torch.relu(x[1]) for x in cnet_list]
|
||||
inp_list = [self.cam(x) * x for x in inp_list]
|
||||
att = [self.sam(x) for x in inp_list]
|
||||
|
||||
geo_fn = Combined_Geo_Encoding_Volume(features_left_04.to(self.dtype), features_right_04.to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels)
|
||||
b, c, h, w = features_left[0].shape
|
||||
coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1)
|
||||
disp = init_disp.to(self.dtype)
|
||||
|
||||
# GRUs iterations to update disparity (1/4 resolution)
|
||||
for itr in range(self.args.valid_iters):
|
||||
disp = disp.detach()
|
||||
geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=True)
|
||||
net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att)
|
||||
|
||||
disp = disp + delta_disp.to(self.dtype)
|
||||
if itr < self.args.valid_iters-1:
|
||||
continue
|
||||
|
||||
disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype))
|
||||
|
||||
return disp_up
|
||||
|
||||
|
||||
class TrtRunner(nn.Module):
|
||||
def __init__(self, args, feature_runner_engine_path, post_runner_engine_path):
|
||||
super().__init__()
|
||||
import tensorrt as trt
|
||||
self.args = args
|
||||
with open(feature_runner_engine_path, 'rb') as file:
|
||||
engine_data = file.read()
|
||||
self.trt_logger = trt.Logger(trt.Logger.WARNING)
|
||||
self.feature_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data)
|
||||
self.feature_context = self.feature_engine.create_execution_context()
|
||||
|
||||
with open(post_runner_engine_path, 'rb') as file:
|
||||
engine_data = file.read()
|
||||
self.post_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data)
|
||||
self.post_context = self.post_engine.create_execution_context()
|
||||
self.max_disp = args.max_disp
|
||||
self.cv_group = args.get('cv_group', 8)
|
||||
|
||||
def trt_dtype_to_torch(self, dt):
|
||||
import tensorrt as trt
|
||||
if dt==trt.DataType.FLOAT: return torch.float32
|
||||
if dt==trt.DataType.HALF: return torch.float16
|
||||
if dt==trt.DataType.BF16: return torch.bfloat16
|
||||
if dt==trt.DataType.INT32: return torch.int32
|
||||
if dt==trt.DataType.INT8: return torch.int8
|
||||
if dt==trt.DataType.BOOL: return torch.bool
|
||||
raise RuntimeError(f"Unsupported TRT dtype: {dt}")
|
||||
|
||||
def get_io_tensor_names(self, engine, mode):
|
||||
names = []
|
||||
n = engine.num_io_tensors
|
||||
for i in range(n):
|
||||
name = engine.get_tensor_name(i)
|
||||
if engine.get_tensor_mode(name)==mode:
|
||||
names.append(name)
|
||||
return names
|
||||
|
||||
def run_trt(self, engine, context, inputs_by_name:dict):
|
||||
import tensorrt as trt
|
||||
for name, tensor in list(inputs_by_name.items()):
|
||||
expected_dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name))
|
||||
if tensor.dtype != expected_dtype: inputs_by_name[name] = tensor.to(expected_dtype)
|
||||
if not inputs_by_name[name].is_contiguous(): inputs_by_name[name] = inputs_by_name[name].contiguous()
|
||||
context.set_input_shape(name, tuple(inputs_by_name[name].shape))
|
||||
outputs = {}
|
||||
out_names = [n for n in self.get_io_tensor_names(engine, trt.TensorIOMode.OUTPUT)]
|
||||
for name in out_names:
|
||||
shp = tuple(context.get_tensor_shape(name))
|
||||
dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name))
|
||||
outputs[name] = torch.empty(shp, device='cuda', dtype=dtype)
|
||||
for name, tensor in inputs_by_name.items(): context.set_tensor_address(name, int(tensor.data_ptr()))
|
||||
for name, tensor in outputs.items(): context.set_tensor_address(name, int(tensor.data_ptr()))
|
||||
stream = torch.cuda.current_stream().cuda_stream
|
||||
ok = context.execute_async_v3(stream)
|
||||
assert ok
|
||||
return outputs
|
||||
|
||||
def forward(self, image1, image2):
|
||||
import tensorrt as trt
|
||||
feat_out = self.run_trt(self.feature_engine, self.feature_context, {'left': image1, 'right': image2})
|
||||
gwc_volume = build_gwc_volume_triton(feat_out['features_left_04'].half(), feat_out['features_right_04'].half(), self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
|
||||
post_inputs = feat_out
|
||||
post_inputs['gwc_volume'] = gwc_volume
|
||||
in_names = self.get_io_tensor_names(self.post_engine, trt.TensorIOMode.INPUT)
|
||||
tmp_keys = list(post_inputs.keys())
|
||||
for k in tmp_keys:
|
||||
if k not in in_names:
|
||||
del post_inputs[k]
|
||||
out = self.run_trt(self.post_engine, self.post_context, post_inputs)
|
||||
disp = out['disp']
|
||||
return disp
|
||||
Reference in New Issue
Block a user