dasha_f
6e1a22ba8b
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>
445 lines
21 KiB
Python
Executable File
445 lines
21 KiB
Python
Executable File
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 |