chenxuanhong
39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
#############################################################
|
|
# File: ResBlock.py
|
|
# Created Date: Monday July 5th 2021
|
|
# Author: Chen Xuanhong
|
|
# Email: chenxuanhongzju@outlook.com
|
|
# Last Modified: Monday, 5th July 2021 12:18:18 am
|
|
# Modified By: Chen Xuanhong
|
|
# Copyright (c) 2021 Shanghai Jiao Tong University
|
|
#############################################################
|
|
|
|
from torch import nn
|
|
|
|
class ResBlock(nn.Module):
|
|
def __init__(self, in_channel, k_size = 3, stride=1):
|
|
super().__init__()
|
|
padding_size = int((k_size -1)/2)
|
|
self.block = nn.Sequential(
|
|
nn.ReflectionPad2d(padding_size),
|
|
nn.Conv2d(in_channels = in_channel , out_channels = in_channel , kernel_size= k_size, stride=stride, bias= False),
|
|
nn.InstanceNorm2d(in_channel, affine=True, momentum=0),
|
|
nn.ReflectionPad2d(padding_size),
|
|
nn.Conv2d(in_channels = in_channel , out_channels = in_channel , kernel_size= k_size, stride=stride, bias= False),
|
|
nn.InstanceNorm2d(in_channel, affine=True, momentum=0)
|
|
)
|
|
self.__weights_init__()
|
|
|
|
def __weights_init__(self):
|
|
for m in self.modules():
|
|
if isinstance(m,nn.Conv2d):
|
|
nn.init.xavier_uniform_(m.weight)
|
|
|
|
def forward(self, input):
|
|
res = input
|
|
h = self.block(input)
|
|
out = h + res
|
|
return out
|