diff --git a/fuse_conv_and_bn.py b/fuse_conv_and_bn.py new file mode 100644 index 0000000..17a31e8 --- /dev/null +++ b/fuse_conv_and_bn.py @@ -0,0 +1,31 @@ +import torch + + +def fuse_conv_and_bn(conv, bn): + """Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.""" + fusedconv = ( + torch.nn.Conv2d( + conv.in_channels, + conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + padding=conv.padding, + dilation=conv.dilation, + groups=conv.groups, + bias=True, + ) + .requires_grad_(False) + .to(conv.weight.device) + ) + + # Prepare filters + w_conv = conv.weight.clone().view(conv.out_channels, -1) + w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) + fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape)) + + # Prepare spatial bias + b_conv = torch.zeros(conv.weight.shape[0], device=conv.weight.device) if conv.bias is None else conv.bias + b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) + fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) + + return fusedconv diff --git a/infer.py b/infer.py index 95fb180..10bede2 100644 --- a/infer.py +++ b/infer.py @@ -4,8 +4,10 @@ import numpy as np import cv2 from NeuFlow.neuflow import NeuFlow + from NeuFlow.backbone_v7 import ConvBlock from data_utils import flow_viz +from fuse_conv_and_bn import fuse_conv_and_bn image_width = 768 @@ -20,80 +22,66 @@ def get_cuda_image(image_path): return image[None].cuda() -def fuse_conv_and_bn(conv, bn): - """Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.""" - fusedconv = ( - torch.nn.Conv2d( - conv.in_channels, - conv.out_channels, - kernel_size=conv.kernel_size, - stride=conv.stride, - padding=conv.padding, - dilation=conv.dilation, - groups=conv.groups, - bias=True, - ) - .requires_grad_(False) - .to(conv.weight.device) - ) - - # Prepare filters - w_conv = conv.weight.clone().view(conv.out_channels, -1) - w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) - fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape)) - - # Prepare spatial bias - b_conv = torch.zeros(conv.weight.shape[0], device=conv.weight.device) if conv.bias is None else conv.bias - b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) - fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) - - return fusedconv +def main(hugging_face: bool) -> None: + image_path_list = sorted(glob('test_images/*.jpg')) + vis_path = 'test_results/' + device = torch.device('cuda') -image_path_list = sorted(glob('test_images/*.jpg')) -vis_path = 'test_results/' + if hugging_face: + print("from_pretrained: Study-is-happy/neuflow-v2") + model = NeuFlow.from_pretrained("Study-is-happy/neuflow-v2").to(device) + else: + model = NeuFlow().to(device) + print("load: neuflow_mixed.pth") + checkpoint = torch.load('neuflow_mixed.pth', map_location='cuda') + model.load_state_dict(checkpoint['model'], strict=True) -device = torch.device('cuda') + for m in model.modules(): + if type(m) is ConvBlock: + m.conv1 = fuse_conv_and_bn(m.conv1, m.norm1) # update conv + m.conv2 = fuse_conv_and_bn(m.conv2, m.norm2) # update conv + delattr(m, "norm1") # remove batchnorm + delattr(m, "norm2") # remove batchnorm + m.forward = m.forward_fuse # update forward -model = NeuFlow().to(device) + model.eval() + model.half() -checkpoint = torch.load('neuflow_mixed.pth', map_location='cuda') + model.init_bhwd(1, image_height, image_width, 'cuda') -model.load_state_dict(checkpoint['model'], strict=True) + if not os.path.exists(vis_path): + os.makedirs(vis_path) -for m in model.modules(): - if type(m) is ConvBlock: - m.conv1 = fuse_conv_and_bn(m.conv1, m.norm1) # update conv - m.conv2 = fuse_conv_and_bn(m.conv2, m.norm2) # update conv - delattr(m, "norm1") # remove batchnorm - delattr(m, "norm2") # remove batchnorm - m.forward = m.forward_fuse # update forward + for image_path_0, image_path_1 in zip(image_path_list[:-1], image_path_list[1:]): -model.eval() -model.half() + print(image_path_0) -model.init_bhwd(1, image_height, image_width, 'cuda') + image_0 = get_cuda_image(image_path_0) + image_1 = get_cuda_image(image_path_1) -if not os.path.exists(vis_path): - os.makedirs(vis_path) + file_name = os.path.basename(image_path_0) -for image_path_0, image_path_1 in zip(image_path_list[:-1], image_path_list[1:]): + with torch.no_grad(): - print(image_path_0) + flow = model(image_0, image_1)[-1][0] - image_0 = get_cuda_image(image_path_0) - image_1 = get_cuda_image(image_path_1) + flow = flow.permute(1,2,0).cpu().numpy() - file_name = os.path.basename(image_path_0) + flow = flow_viz.flow_to_image(flow) - with torch.no_grad(): + image_0 = cv2.resize(cv2.imread(image_path_0), (image_width, image_height)) - flow = model(image_0, image_1)[-1][0] + cv2.imwrite(vis_path + file_name, np.vstack([image_0, flow])) - flow = flow.permute(1,2,0).cpu().numpy() - - flow = flow_viz.flow_to_image(flow) - image_0 = cv2.resize(cv2.imread(image_path_0), (image_width, image_height)) +if __name__ == "__main__": + from argparse import ArgumentParser - cv2.imwrite(vis_path + file_name, np.vstack([image_0, flow])) + parser = ArgumentParser() + parser.add_argument( + "--hugging-face", action="store_true", + help="load model form hugging face (Study-is-happy/neuflow-v2)" + ) + args = parser.parse_args() + main(args.hugging_face) diff --git a/infer_hf.py b/infer_hf.py deleted file mode 100644 index be30904..0000000 --- a/infer_hf.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -from glob import glob -import os -import numpy as np -import cv2 -from NeuFlow.neuflow import NeuFlow -from NeuFlow.backbone_v7 import ConvBlock -from data_utils import flow_viz - - -image_width = 768 -image_height = 432 - -def get_cuda_image(image_path): - image = cv2.imread(image_path) - - image = cv2.resize(image, (image_width, image_height)) - - image = torch.from_numpy(image).permute(2, 0, 1).half() - return image[None].cuda() - - -def fuse_conv_and_bn(conv, bn): - """Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.""" - fusedconv = ( - torch.nn.Conv2d( - conv.in_channels, - conv.out_channels, - kernel_size=conv.kernel_size, - stride=conv.stride, - padding=conv.padding, - dilation=conv.dilation, - groups=conv.groups, - bias=True, - ) - .requires_grad_(False) - .to(conv.weight.device) - ) - - # Prepare filters - w_conv = conv.weight.clone().view(conv.out_channels, -1) - w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) - fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape)) - - # Prepare spatial bias - b_conv = torch.zeros(conv.weight.shape[0], device=conv.weight.device) if conv.bias is None else conv.bias - b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) - fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) - - return fusedconv - - -image_path_list = sorted(glob('test_images/*.jpg')) -vis_path = 'test_results/' - -device = torch.device('cuda') - -model = NeuFlow.from_pretrained("Study-is-happy/neuflow-v2").to(device) - -for m in model.modules(): - if type(m) is ConvBlock: - m.conv1 = fuse_conv_and_bn(m.conv1, m.norm1) # update conv - m.conv2 = fuse_conv_and_bn(m.conv2, m.norm2) # update conv - delattr(m, "norm1") # remove batchnorm - delattr(m, "norm2") # remove batchnorm - m.forward = m.forward_fuse # update forward - -model.eval() -model.half() - -model.init_bhwd(1, image_height, image_width, 'cuda') - -if not os.path.exists(vis_path): - os.makedirs(vis_path) - -for image_path_0, image_path_1 in zip(image_path_list[:-1], image_path_list[1:]): - - print(image_path_0) - - image_0 = get_cuda_image(image_path_0) - image_1 = get_cuda_image(image_path_1) - - file_name = os.path.basename(image_path_0) - - with torch.no_grad(): - - flow = model(image_0, image_1)[-1][0] - - flow = flow.permute(1,2,0).cpu().numpy() - - flow = flow_viz.flow_to_image(flow) - - image_0 = cv2.resize(cv2.imread(image_path_0), (image_width, image_height)) - - cv2.imwrite(vis_path + file_name, np.vstack([image_0, flow])) -