def __init__(self, maxdisp):
        super(PSMNet, self).__init__()
        self.maxdisp = maxdisp

        self.feature_extraction = submodule.feature_extraction()

        self.dres0 = nn.Sequential(submodule.convbn_3d(64, 32, 3, 1, 1),
                                   nn.ReLU(inplace=True),
                                   submodule.convbn_3d(32, 32, 3, 1, 1),
                                   nn.ReLU(inplace=True))

        self.dres1 = nn.Sequential(submodule.convbn_3d(32, 32, 3, 1, 1),
                                   nn.ReLU(inplace=True),
                                   submodule.convbn_3d(32, 32, 3, 1, 1))

        self.dres2 = hourglass(32)

        self.dres3 = hourglass(32)

        self.dres4 = hourglass(32)

        self.classif1 = nn.Sequential(
            submodule.convbn_3d(32, 32, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False))

        self.classif2 = nn.Sequential(
            submodule.convbn_3d(32, 32, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False))

        self.classif3 = nn.Sequential(
            submodule.convbn_3d(32, 32, 3, 1, 1), nn.ReLU(inplace=True),
            nn.Conv3d(32, 1, kernel_size=3, padding=1, stride=1, bias=False))

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.Conv3d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[
                    2] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()
    def __init__(self, maxdisp):
        super(PSMNet, self).__init__()
        self.maxdisp = maxdisp
        self.inplanes = 32
        self.feature_extraction = sb.feature_extraction()

        self.dres0 = nn.Sequential(sb.convbn_3d(64, 32, 3, 1, 1),
                                   nn.ReLU(inplace=True))

        self.dres1 = nn.Sequential(sb.convbn_3d(32, 32, 3, 1, 1),
                                   nn.ReLU(inplace=True),
                                   sb.convbn_3d(32, 32, 3, 1, 1))

        self.dres2 = nn.Sequential(
            sb.convbn_3d(self.inplanes,
                         self.inplanes,
                         kernel_size=3,
                         stride=2,
                         pad=1), nn.ReLU(inplace=True),
            sb.convbn_3d(self.inplanes,
                         self.inplanes,
                         kernel_size=3,
                         stride=1,
                         pad=1))

        self.dres3 = nn.Sequential(
            sb.convbn_3d(self.inplanes,
                         self.inplanes * 2,
                         kernel_size=3,
                         stride=2,
                         pad=1), nn.ReLU(inplace=True),
            sb.convbn_3d(self.inplanes * 2,
                         self.inplanes * 2,
                         kernel_size=3,
                         stride=1,
                         pad=1))

        self.dres4 = nn.Sequential(
            sb.convbn_3d(self.inplanes * 2,
                         self.inplanes * 2,
                         kernel_size=3,
                         stride=2,
                         pad=1), nn.ReLU(inplace=True),
            sb.convbn_3d(self.inplanes * 2,
                         self.inplanes * 2,
                         kernel_size=3,
                         stride=1,
                         pad=1))

        self.classif1 = nn.Sequential(sb.convbn_3d(128, 64, 3, 1, 1),
                                      nn.ReLU(inplace=True),
                                      nn.Conv3d(64, 32, 3, 1, 1))
        self.classif2 = nn.Sequential(sb.convbn_3d(128, 64, 3, 1, 1),
                                      nn.ReLU(inplace=True),
                                      nn.Conv3d(64, 32, 3, 1, 1))

        self.conv1_1 = nn.Sequential(sb.convbn_3d(32, 32, 3, 1, 1),
                                     nn.ReLU(inplace=True))

        self.conv2_1 = nn.Sequential(sb.convbn_3d(32, 32, 3, 1, 1),
                                     nn.ReLU(inplace=True))

        self.dres5 = nn.Sequential(
            sb.convbn_3d(self.inplanes * 2,
                         self.inplanes,
                         kernel_size=3,
                         stride=1,
                         pad=1), nn.ReLU(inplace=True))

        self.deconv3_1 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes * 2,
                               self.inplanes * 2,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes * 2))

        self.deconv4_1 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes * 2,
                               self.inplanes,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv5_1 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               self.inplanes,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv6_1 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               self.inplanes,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv3_2 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes * 2,
                               self.inplanes * 2,
                               kernel_size=5,
                               padding=2,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes * 2))

        self.deconv4_2 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes * 2,
                               self.inplanes,
                               kernel_size=5,
                               padding=2,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv5_2 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               self.inplanes,
                               kernel_size=5,
                               padding=2,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv6_2 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               self.inplanes,
                               kernel_size=5,
                               padding=2,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv1 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               self.inplanes,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(self.inplanes))

        self.deconv2 = nn.Sequential(
            nn.ConvTranspose3d(self.inplanes,
                               1,
                               kernel_size=3,
                               padding=1,
                               output_padding=1,
                               stride=2,
                               bias=False), nn.BatchNorm3d(1))

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.Conv3d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[
                    2] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm3d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()