原文链接:Tutorial 4: Adding New Modules — mmflow documentation
教程4 添加新的模块
MMFlow把一个光流估计方法flow_estimator拆分成了编码器encoder和解码器decoder。本教程展示如何添加新的部件。
添加新的编码器
1.创建一个新文件:mmflow/models/encoders/my_model.py
from mmcv.runner import BaseModule
from ..builder import ENCODERS
@ENCODERS.register_module()
class MyModel(BaseModule):
def __init__(self, arg1, arg2):
pass
def forward(self, x): # should return a tuple
pass
def init_weights(self, pretrained=None):
pass
2.在mmflow/models/encoders/__init__.py导入该模块(MyModel)
from .my_model import MyModel
添加新的解码器
1.创建新文件:mmflow/models/decoders/my_decoder.py
你可以写一个新的从MMCV的BaseModule中继承的头,然后覆写forward(self, x), forward_train和forward_test等方法。我们在MMCV中有一个统一的接口用于权重初始化,你可以使用init_cfg用于指定初始化函数和参数,或者覆写init_weights()方法。(如果你更喜欢自定义初始化的话)
from ..builder import DECODERS
@DECODERS.register_module()
class MyDecoder(BaseModule):
def __init__(self, arg1, arg2):
pass
def forward(self, *args):
pass
# optional
def init_weights(self):
pass
def forward_train(self, *args, flow_gt):
flow_pred = self.forward(*args)
return self.losses(flow_pred, flow_gt)
def forward_test(self,*args, img_metas):
flow_pred = self.forward(*args)
return self.get_flow(flow_pred, img_metas)
loss是用于计算模型输出和目标之间损失值的损失函数,get_flow在BaseDecoder中实现,用于把光流的形状尺寸恢复到原始输入图像的尺寸。
1.导入mmflow/models/decoders/__init__.py中的模块
from .my_decoder import MyDecoder
添加新的光流估计器
1.创建新文件:mmflow/models/flow_estimators/my_estimator.py
你可以写一个新的类似于PWC-Net的继承自FlowEstimator的头,然后实现forward_train和forward_test等方法。
from ..builder import FLOW_ESTIMATORS
from .base import FlowEstimator
@FLOW_ESTIMATORS.register_module()
class MyEstimator(FlowEstimator):
def __init__(self, arg1, arg2):
pass
def forward_train(self, imgs):
pass
def forward_test(self, imgs):
pass
2.导入mmflow/models/flow_estimator/__init__.py中的模块
from .my_estimator import MyEstimator
3.在你的配置文件中使用它
我们把模型类型设为MyEstimator
model = dict(
type='MyEstimator',
encoder=dict(
type='MyModel',
arg1=xxx,
arg2=xxx),
decoder=dict(
type='MyDecoder',
arg1=xxx,
arg2=xxx))
添加新的损失函数
假设你想增加一个新的名为MyLoss的损失函数用于光流估计。为了添加新的损失函数,使用者需要在mmflow/models/losses/my_loss.py中编程实现。
import torch
import torch.nn as nn
from mmflow.models import LOSSES
def my_loss(pred, target):
pass
@LOSSES.register_module()
class MyLoss(nn.Module):
def __init__(self, arg1):
super(MyLoss, self).__init__()
def forward(self, output, target):
return my_loss(output, target)
然后把它添加到mmflow/models/losses/__init__.py中
from .my_loss import MyLoss, my_loss
要想使用它,需要在模型的flow_loss字段中修改。
flow_loss=dict(type='MyLoss', use_target_weight=False)