Upload GPReconResNet

0a4a687 verified about 2 years ago

10.5 kB

# Adapted from https://raw.githubusercontent.com/soumickmj/NCC1701/main/Bridge/models/ResNet/MickResNet.py

import torch
import torch.nn as nn
import numpy as np
import sys
import torch.nn.functional as F
from tricorder.torch.transforms import Interpolator

__author__ = "Soumick Chatterjee"
__copyright__ = "Copyright 2019, Soumick Chatterjee & OvGU:ESF:MEMoRIAL"
__credits__ = ["Soumick Chatterjee"]

__license__ = "GPL"
__version__ = "1.0.0"
__email__ = "soumick.chatterjee@ovgu.de"
__status__ = "Published"

class ResidualBlock(nn.Module):
    def __init__(self, in_features, drop_prob=0.2): #drop_prob=0.2
        super(ResidualBlock, self).__init__()

        conv_block = [  layer_pad(1),
                        layer_conv(in_features, in_features, 3),
                        layer_norm(in_features),
                        act_relu(),
                        layer_drop(p=drop_prob, inplace=True),
                        layer_pad(1),
                        layer_conv(in_features, in_features, 3) ,
                        layer_norm(in_features) ]

        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x):
        return x + self.conv_block(x)

class DownsamplingBlock(nn.Module):
    def __init__(self, in_features, out_features):
        super(DownsamplingBlock, self).__init__()

        conv_block = [  layer_conv(in_features, out_features, 3, stride=2, padding=1),
                        layer_norm(out_features),
                        act_relu()  ]
        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x):
        return self.conv_block(x)

class UpsamplingBlock(nn.Module):
    def __init__(self, in_features, out_features, mode="upconv", interpolator=None, post_interp_convtrans=False):
        super(UpsamplingBlock, self).__init__()

        self.interpolator = interpolator
        self.mode = mode
        self.post_interp_convtrans = post_interp_convtrans
        if self.post_interp_convtrans:
            self.post_conv = layer_conv(out_features, out_features, 1)

        if mode == "upconv":
            conv_block = [  layer_convtrans(in_features, out_features, 3, stride=2, padding=1, output_padding=1),   ]
        else:
            conv_block = [  layer_pad(1),
                            layer_conv(in_features, out_features, 3),   ]
        conv_block += [ layer_norm(out_features),
                        act_relu()  ]
        self.conv_block = nn.Sequential(*conv_block)

    def forward(self, x, out_shape=None):
        if self.mode == "upconv":
            if self.post_interp_convtrans:
                x = self.conv_block(x)
                if x.shape[2:] != out_shape:
                    return self.post_conv(self.interpolator(x, out_shape))
                else:
                    return x
            else:
                return self.conv_block(x)
        else:
            return self.conv_block(self.interpolator(x, out_shape))

class GP_ReconResNet(nn.Module):
    def __init__(self, in_channels=1, n_classes=1, res_blocks=14, starting_nfeatures=64, updown_blocks=2, is_relu_leaky=True, do_batchnorm=False, res_drop_prob=0.2,    #res_drop_prob=0.2
                        out_act="softmax", forwardV=0, upinterp_algo='upconv', post_interp_convtrans=False, is3D=False): #should use 14 as that gives number of trainable parameters close to number of possible pixel values in a image 256x256 
        super(GP_ReconResNet, self).__init__()

        layers = {}
        if is3D:
            sys.exit("ResNet: for implemented for 3D, ReflectionPad3d code is required")
            layers["layer_conv"] = nn.Conv3d
            layers["layer_convtrans"] = nn.ConvTranspose3d
            if do_batchnorm:
                layers["layer_norm"] = nn.BatchNorm3d
            else:
                layers["layer_norm"] = nn.InstanceNorm3d
            layers["layer_drop"] = nn.Dropout3d
            layers["layer_pad"] = ReflectionPad3d
            layers["interp_mode"] = 'trilinear'
        else:
            layers["layer_conv"] = nn.Conv2d
            layers["layer_convtrans"] = nn.ConvTranspose2d
            if do_batchnorm:
                layers["layer_norm"] = nn.BatchNorm2d
            else:
                layers["layer_norm"] = nn.InstanceNorm2d
            layers["layer_drop"] = nn.Dropout2d
            layers["layer_pad"] = nn.ReflectionPad2d
            layers["interp_mode"] = 'bilinear'            
        if is_relu_leaky:
            layers["act_relu"] = nn.PReLU
        else:
            layers["act_relu"] = nn.ReLU
        globals().update(layers)

        self.forwardV = forwardV
        self.upinterp_algo = upinterp_algo

        interpolator = Interpolator(mode=layers["interp_mode"] if self.upinterp_algo == "upconv" else self.upinterp_algo)

        in_channels = in_channels
        out_channels = n_classes
        # Initial convolution block
        intialConv = [  layer_pad(3),
                        layer_conv(in_channels, starting_nfeatures, 7),
                        layer_norm(starting_nfeatures),
                        act_relu() ]

        # Downsampling [need to save the shape for upsample]
        downsam = []
        in_features = starting_nfeatures
        out_features = in_features*2
        for _ in range(updown_blocks):
            downsam.append(DownsamplingBlock(in_features, out_features))
            in_features = out_features
            out_features = in_features*2

        # Residual blocks
        resblocks = []
        for _ in range(res_blocks):
            resblocks += [ResidualBlock(in_features, res_drop_prob)]

        # Upsampling
        upsam = []
        out_features = in_features//2
        for _ in range(updown_blocks):
            upsam.append(UpsamplingBlock(in_features, out_features, self.upinterp_algo, interpolator, post_interp_convtrans))
            in_features = out_features
            out_features = in_features//2

        # Output layer
        finalconv = [   layer_conv(starting_nfeatures, out_channels, 1),    ] #kernel size changed from 7 to 1 to make GMP work

        if out_act == "sigmoid":
            finalconv += [   nn.Sigmoid(), ]
        elif out_act == "relu":
            finalconv += [   act_relu(), ]
        elif out_act == "tanh":
            finalconv += [   nn.Tanh(), ]
        elif out_act == "softmax":
            finalconv += [   nn.Softmax2d(), ]


        self.intialConv = nn.Sequential(*intialConv)
        self.downsam = nn.ModuleList(downsam)
        self.resblocks = nn.Sequential(*resblocks)
        self.upsam = nn.ModuleList(upsam)
        self.finalconv = nn.Sequential(*finalconv)

        ### For Classification, following Florian's GP-UNet
        self.GMP = nn.AdaptiveMaxPool2d((1, 1))

        if self.forwardV == 0:
            self.forward = self.forwardV0
        elif self.forwardV == 1:
            sys.exit("ResNet: its identical to V0 in case of GP_ResNet")
        elif self.forwardV == 2:
            self.forward = self.forwardV2
        elif self.forwardV == 3:
            self.forward = self.forwardV3
        elif self.forwardV == 4:
            self.forward = self.forwardV4
        elif self.forwardV == 5:
            self.forward = self.forwardV5

    def final_step(self, x):    
        if self.training:
            x = self.GMP(x)
            return self.finalconv(x).view(x.shape[0],-1)
        else:
            mask = self.finalconv(x)
            x = self.GMP(x)
            pred = self.finalconv(x).view(x.shape[0],-1)
            return pred, mask
            
    def forwardV0(self, x):
        #v0: Original Version
        x = self.intialConv(x)
        shapes = []
        for downblock in self.downsam:
            shapes.append(x.shape[2:])
            x = downblock(x)
        x = self.resblocks(x)
        for i, upblock in enumerate(self.upsam):
            x = upblock(x, shapes[-1-i])
        return self.final_step(x)

    def forwardV2(self, x):
        #v2: residual of v1 + input to the residual blocks added back with the output
        out = self.intialConv(x)
        shapes = []
        for downblock in self.downsam:
            shapes.append(out.shape[2:])
            out = downblock(out)
        out = out + self.resblocks(out)
        for i, upblock in enumerate(self.upsam):
            out = upblock(out, shapes[-1-i])
        return self.final_step(out)

    def forwardV3(self, x):
        #v3: residual of v2 + input of the initial conv added back with the output
        out = x + self.intialConv(x)
        shapes = []
        for downblock in self.downsam:
            shapes.append(out.shape[2:])
            out = downblock(out)
        out = out + self.resblocks(out)
        for i, upblock in enumerate(self.upsam):
            out = upblock(out, shapes[-1-i])
        return self.final_step(out)

    def forwardV4(self, x):
        #v4: residual of v3 + output of the initial conv added back with the input of final conv
        iniconv = x + self.intialConv(x)
        shapes = []
        if len(self.downsam) > 0:
            for i, downblock in enumerate(self.downsam):
                if i == 0:
                    shapes.append(iniconv.shape[2:])
                    out = downblock(iniconv)
                else:
                    shapes.append(out.shape[2:])
                    out = downblock(out)
        else:
            out = iniconv
        out = out + self.resblocks(out)
        for i, upblock in enumerate(self.upsam):
            out = upblock(out, shapes[-1-i])
        out = iniconv + out
        return self.final_step(out)

    def forwardV5(self, x):
        #v5: residual of v4 + individual down blocks with individual up blocks
        outs = [x + self.intialConv(x)]
        shapes = []
        for i, downblock in enumerate(self.downsam):
            shapes.append(outs[-1].shape[2:])
            outs.append(downblock(outs[-1]))
        outs[-1] = outs[-1] + self.resblocks(outs[-1])
        for i, upblock in enumerate(self.upsam):
            outs[-1] = upblock(outs[-1], shapes[-1-i])
            outs[-1] = outs[-2] + outs.pop()
        return self.final_step(outs.pop())

#to run it here from this script, uncomment the following

if __name__ == "__main__":                #to run it
    image = torch.rand(2, 1, 240, 240)    #specify your image: batch size, Channel, height, width
    model = GP_ReconResNet(in_channels=1, n_classes=3, upinterp_algo='sinc')                        #Initialize the model
    # model.eval()
    out = model(image)
    print(model(image))

	# Adapted from https://raw.githubusercontent.com/soumickmj/NCC1701/main/Bridge/models/ResNet/MickResNet.py

	import torch
	import torch.nn as nn
	import numpy as np
	import sys
	import torch.nn.functional as F
	from tricorder.torch.transforms import Interpolator

	__author__ = "Soumick Chatterjee"
	__copyright__ = "Copyright 2019, Soumick Chatterjee & OvGU:ESF:MEMoRIAL"
	__credits__ = ["Soumick Chatterjee"]

	__license__ = "GPL"
	__version__ = "1.0.0"
	__email__ = "soumick.chatterjee@ovgu.de"
	__status__ = "Published"

	class ResidualBlock(nn.Module):
	def __init__(self, in_features, drop_prob=0.2): #drop_prob=0.2
	super(ResidualBlock, self).__init__()

	conv_block = [ layer_pad(1),
	layer_conv(in_features, in_features, 3),
	layer_norm(in_features),
	act_relu(),
	layer_drop(p=drop_prob, inplace=True),
	layer_pad(1),
	layer_conv(in_features, in_features, 3) ,
	layer_norm(in_features) ]

	self.conv_block = nn.Sequential(*conv_block)

	def forward(self, x):
	return x + self.conv_block(x)

	class DownsamplingBlock(nn.Module):
	def __init__(self, in_features, out_features):
	super(DownsamplingBlock, self).__init__()

	conv_block = [ layer_conv(in_features, out_features, 3, stride=2, padding=1),
	layer_norm(out_features),
	act_relu() ]
	self.conv_block = nn.Sequential(*conv_block)

	def forward(self, x):
	return self.conv_block(x)

	class UpsamplingBlock(nn.Module):
	def __init__(self, in_features, out_features, mode="upconv", interpolator=None, post_interp_convtrans=False):
	super(UpsamplingBlock, self).__init__()

	self.interpolator = interpolator
	self.mode = mode
	self.post_interp_convtrans = post_interp_convtrans
	if self.post_interp_convtrans:
	self.post_conv = layer_conv(out_features, out_features, 1)

	if mode == "upconv":
	conv_block = [ layer_convtrans(in_features, out_features, 3, stride=2, padding=1, output_padding=1), ]
	else:
	conv_block = [ layer_pad(1),
	layer_conv(in_features, out_features, 3), ]
	conv_block += [ layer_norm(out_features),
	act_relu() ]
	self.conv_block = nn.Sequential(*conv_block)

	def forward(self, x, out_shape=None):
	if self.mode == "upconv":
	if self.post_interp_convtrans:
	x = self.conv_block(x)
	if x.shape[2:] != out_shape:
	return self.post_conv(self.interpolator(x, out_shape))
	else:
	return x
	else:
	return self.conv_block(x)
	else:
	return self.conv_block(self.interpolator(x, out_shape))

	class GP_ReconResNet(nn.Module):
	def __init__(self, in_channels=1, n_classes=1, res_blocks=14, starting_nfeatures=64, updown_blocks=2, is_relu_leaky=True, do_batchnorm=False, res_drop_prob=0.2, #res_drop_prob=0.2
	out_act="softmax", forwardV=0, upinterp_algo='upconv', post_interp_convtrans=False, is3D=False): #should use 14 as that gives number of trainable parameters close to number of possible pixel values in a image 256x256
	super(GP_ReconResNet, self).__init__()

	layers = {}
	if is3D:
	sys.exit("ResNet: for implemented for 3D, ReflectionPad3d code is required")
	layers["layer_conv"] = nn.Conv3d
	layers["layer_convtrans"] = nn.ConvTranspose3d
	if do_batchnorm:
	layers["layer_norm"] = nn.BatchNorm3d
	else:
	layers["layer_norm"] = nn.InstanceNorm3d
	layers["layer_drop"] = nn.Dropout3d
	layers["layer_pad"] = ReflectionPad3d
	layers["interp_mode"] = 'trilinear'
	else:
	layers["layer_conv"] = nn.Conv2d
	layers["layer_convtrans"] = nn.ConvTranspose2d
	if do_batchnorm:
	layers["layer_norm"] = nn.BatchNorm2d
	else:
	layers["layer_norm"] = nn.InstanceNorm2d
	layers["layer_drop"] = nn.Dropout2d
	layers["layer_pad"] = nn.ReflectionPad2d
	layers["interp_mode"] = 'bilinear'
	if is_relu_leaky:
	layers["act_relu"] = nn.PReLU
	else:
	layers["act_relu"] = nn.ReLU
	globals().update(layers)

	self.forwardV = forwardV
	self.upinterp_algo = upinterp_algo

	interpolator = Interpolator(mode=layers["interp_mode"] if self.upinterp_algo == "upconv" else self.upinterp_algo)

	in_channels = in_channels
	out_channels = n_classes
	# Initial convolution block
	intialConv = [ layer_pad(3),
	layer_conv(in_channels, starting_nfeatures, 7),
	layer_norm(starting_nfeatures),
	act_relu() ]

	# Downsampling [need to save the shape for upsample]
	downsam = []
	in_features = starting_nfeatures
	out_features = in_features*2
	for _ in range(updown_blocks):
	downsam.append(DownsamplingBlock(in_features, out_features))
	in_features = out_features
	out_features = in_features*2

	# Residual blocks
	resblocks = []
	for _ in range(res_blocks):
	resblocks += [ResidualBlock(in_features, res_drop_prob)]

	# Upsampling
	upsam = []
	out_features = in_features//2
	for _ in range(updown_blocks):
	upsam.append(UpsamplingBlock(in_features, out_features, self.upinterp_algo, interpolator, post_interp_convtrans))
	in_features = out_features
	out_features = in_features//2

	# Output layer
	finalconv = [ layer_conv(starting_nfeatures, out_channels, 1), ] #kernel size changed from 7 to 1 to make GMP work

	if out_act == "sigmoid":
	finalconv += [ nn.Sigmoid(), ]
	elif out_act == "relu":
	finalconv += [ act_relu(), ]
	elif out_act == "tanh":
	finalconv += [ nn.Tanh(), ]
	elif out_act == "softmax":
	finalconv += [ nn.Softmax2d(), ]


	self.intialConv = nn.Sequential(*intialConv)
	self.downsam = nn.ModuleList(downsam)
	self.resblocks = nn.Sequential(*resblocks)
	self.upsam = nn.ModuleList(upsam)
	self.finalconv = nn.Sequential(*finalconv)

	### For Classification, following Florian's GP-UNet
	self.GMP = nn.AdaptiveMaxPool2d((1, 1))

	if self.forwardV == 0:
	self.forward = self.forwardV0
	elif self.forwardV == 1:
	sys.exit("ResNet: its identical to V0 in case of GP_ResNet")
	elif self.forwardV == 2:
	self.forward = self.forwardV2
	elif self.forwardV == 3:
	self.forward = self.forwardV3
	elif self.forwardV == 4:
	self.forward = self.forwardV4
	elif self.forwardV == 5:
	self.forward = self.forwardV5

	def final_step(self, x):
	if self.training:
	x = self.GMP(x)
	return self.finalconv(x).view(x.shape[0],-1)
	else:
	mask = self.finalconv(x)
	x = self.GMP(x)
	pred = self.finalconv(x).view(x.shape[0],-1)
	return pred, mask

	def forwardV0(self, x):
	#v0: Original Version
	x = self.intialConv(x)
	shapes = []
	for downblock in self.downsam:
	shapes.append(x.shape[2:])
	x = downblock(x)
	x = self.resblocks(x)
	for i, upblock in enumerate(self.upsam):
	x = upblock(x, shapes[-1-i])
	return self.final_step(x)

	def forwardV2(self, x):
	#v2: residual of v1 + input to the residual blocks added back with the output
	out = self.intialConv(x)
	shapes = []
	for downblock in self.downsam:
	shapes.append(out.shape[2:])
	out = downblock(out)
	out = out + self.resblocks(out)
	for i, upblock in enumerate(self.upsam):
	out = upblock(out, shapes[-1-i])
	return self.final_step(out)

	def forwardV3(self, x):
	#v3: residual of v2 + input of the initial conv added back with the output
	out = x + self.intialConv(x)
	shapes = []
	for downblock in self.downsam:
	shapes.append(out.shape[2:])
	out = downblock(out)
	out = out + self.resblocks(out)
	for i, upblock in enumerate(self.upsam):
	out = upblock(out, shapes[-1-i])
	return self.final_step(out)

	def forwardV4(self, x):
	#v4: residual of v3 + output of the initial conv added back with the input of final conv
	iniconv = x + self.intialConv(x)
	shapes = []
	if len(self.downsam) > 0:
	for i, downblock in enumerate(self.downsam):
	if i == 0:
	shapes.append(iniconv.shape[2:])
	out = downblock(iniconv)
	else:
	shapes.append(out.shape[2:])
	out = downblock(out)
	else:
	out = iniconv
	out = out + self.resblocks(out)
	for i, upblock in enumerate(self.upsam):
	out = upblock(out, shapes[-1-i])
	out = iniconv + out
	return self.final_step(out)

	def forwardV5(self, x):
	#v5: residual of v4 + individual down blocks with individual up blocks
	outs = [x + self.intialConv(x)]
	shapes = []
	for i, downblock in enumerate(self.downsam):
	shapes.append(outs[-1].shape[2:])
	outs.append(downblock(outs[-1]))
	outs[-1] = outs[-1] + self.resblocks(outs[-1])
	for i, upblock in enumerate(self.upsam):
	outs[-1] = upblock(outs[-1], shapes[-1-i])
	outs[-1] = outs[-2] + outs.pop()
	return self.final_step(outs.pop())

	#to run it here from this script, uncomment the following

	if __name__ == "__main__": #to run it
	image = torch.rand(2, 1, 240, 240) #specify your image: batch size, Channel, height, width
	model = GP_ReconResNet(in_channels=1, n_classes=3, upinterp_algo='sinc') #Initialize the model
	# model.eval()
	out = model(image)
	print(model(image))