Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
| # holder of all proprietary rights on this computer program. | |
| # You can only use this computer program if you have closed | |
| # a license agreement with MPG or you get the right to use the computer | |
| # program from someone who is authorized to grant you that right. | |
| # Any use of the computer program without a valid license is prohibited and | |
| # liable to prosecution. | |
| # | |
| # Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
| # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
| # for Intelligent Systems. All rights reserved. | |
| # | |
| # Contact: [email protected] | |
| import torch.nn as nn | |
| import pytorch_lightning as pl | |
| class BaseNetwork(pl.LightningModule): | |
| def __init__(self): | |
| super(BaseNetwork, self).__init__() | |
| def init_weights(self, init_type='xavier', gain=0.02): | |
| ''' | |
| initializes network's weights | |
| init_type: normal | xavier | kaiming | orthogonal | |
| https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39 | |
| ''' | |
| def init_func(m): | |
| classname = m.__class__.__name__ | |
| if hasattr(m, 'weight') and (classname.find('Conv') != -1 | |
| or classname.find('Linear') != -1): | |
| if init_type == 'normal': | |
| nn.init.normal_(m.weight.data, 0.0, gain) | |
| elif init_type == 'xavier': | |
| nn.init.xavier_normal_(m.weight.data, gain=gain) | |
| elif init_type == 'kaiming': | |
| nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') | |
| elif init_type == 'orthogonal': | |
| nn.init.orthogonal_(m.weight.data, gain=gain) | |
| if hasattr(m, 'bias') and m.bias is not None: | |
| nn.init.constant_(m.bias.data, 0.0) | |
| elif classname.find('BatchNorm2d') != -1: | |
| nn.init.normal_(m.weight.data, 1.0, gain) | |
| nn.init.constant_(m.bias.data, 0.0) | |
| self.apply(init_func) | |
| class Residual3D(BaseNetwork): | |
| def __init__(self, numIn, numOut): | |
| super(Residual3D, self).__init__() | |
| self.numIn = numIn | |
| self.numOut = numOut | |
| self.with_bias = True | |
| # self.bn = nn.GroupNorm(4, self.numIn) | |
| self.bn = nn.BatchNorm3d(self.numIn) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv1 = nn.Conv3d(self.numIn, | |
| self.numOut, | |
| bias=self.with_bias, | |
| kernel_size=3, | |
| stride=1, | |
| padding=2, | |
| dilation=2) | |
| # self.bn1 = nn.GroupNorm(4, self.numOut) | |
| self.bn1 = nn.BatchNorm3d(self.numOut) | |
| self.conv2 = nn.Conv3d(self.numOut, | |
| self.numOut, | |
| bias=self.with_bias, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1) | |
| # self.bn2 = nn.GroupNorm(4, self.numOut) | |
| self.bn2 = nn.BatchNorm3d(self.numOut) | |
| self.conv3 = nn.Conv3d(self.numOut, | |
| self.numOut, | |
| bias=self.with_bias, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1) | |
| if self.numIn != self.numOut: | |
| self.conv4 = nn.Conv3d(self.numIn, | |
| self.numOut, | |
| bias=self.with_bias, | |
| kernel_size=1) | |
| self.init_weights() | |
| def forward(self, x): | |
| residual = x | |
| # out = self.bn(x) | |
| # out = self.relu(out) | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = self.conv2(out) | |
| out = self.bn2(out) | |
| # out = self.conv3(out) | |
| # out = self.relu(out) | |
| if self.numIn != self.numOut: | |
| residual = self.conv4(x) | |
| return out + residual | |
| class VolumeEncoder(BaseNetwork): | |
| """CycleGan Encoder""" | |
| def __init__(self, num_in=3, num_out=32, num_stacks=2): | |
| super(VolumeEncoder, self).__init__() | |
| self.num_in = num_in | |
| self.num_out = num_out | |
| self.num_inter = 8 | |
| self.num_stacks = num_stacks | |
| self.with_bias = True | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv1 = nn.Conv3d(self.num_in, | |
| self.num_inter, | |
| bias=self.with_bias, | |
| kernel_size=5, | |
| stride=2, | |
| padding=4, | |
| dilation=2) | |
| # self.bn1 = nn.GroupNorm(4, self.num_inter) | |
| self.bn1 = nn.BatchNorm3d(self.num_inter) | |
| self.conv2 = nn.Conv3d(self.num_inter, | |
| self.num_out, | |
| bias=self.with_bias, | |
| kernel_size=5, | |
| stride=2, | |
| padding=4, | |
| dilation=2) | |
| # self.bn2 = nn.GroupNorm(4, self.num_out) | |
| self.bn2 = nn.BatchNorm3d(self.num_out) | |
| self.conv_out1 = nn.Conv3d(self.num_out, | |
| self.num_out, | |
| bias=self.with_bias, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1, | |
| dilation=1) | |
| self.conv_out2 = nn.Conv3d(self.num_out, | |
| self.num_out, | |
| bias=self.with_bias, | |
| kernel_size=3, | |
| stride=1, | |
| padding=1, | |
| dilation=1) | |
| for idx in range(self.num_stacks): | |
| self.add_module("res" + str(idx), | |
| Residual3D(self.num_out, self.num_out)) | |
| self.init_weights() | |
| def forward(self, x, intermediate_output=True): | |
| out = self.conv1(x) | |
| out = self.bn1(out) | |
| out = self.relu(out) | |
| out = self.conv2(out) | |
| out = self.bn2(out) | |
| out = self.relu(out) | |
| out_lst = [] | |
| for idx in range(self.num_stacks): | |
| out = self._modules["res" + str(idx)](out) | |
| out_lst.append(out) | |
| if intermediate_output: | |
| return out_lst | |
| else: | |
| return [out_lst[-1]] | |