Spaces:
Running
Running
| import torch | |
| from torch import nn | |
| def conv4x4(in_c, out_c): | |
| return nn.Sequential( | |
| nn.Conv2d(in_c, out_c,kernel_size=4, stride=2, padding=1), | |
| nn.BatchNorm2d(out_c), | |
| nn.LeakyReLU(0.1, inplace=True), | |
| ) | |
| def deconv4x4(in_c, out_c): | |
| return nn.Sequential( | |
| nn.ConvTranspose2d(in_c, out_c, kernel_size=4, stride=2, padding=1), | |
| nn.BatchNorm2d(out_c), | |
| nn.LeakyReLU(0.1, inplace=True), | |
| ) | |
| class Hear_Net(nn.Module): | |
| def __init__(self): | |
| super(Hear_Net, self).__init__() | |
| self.down1 = conv4x4(6, 64) | |
| self.down2 = conv4x4(64, 128) | |
| self.down3 = conv4x4(128, 256) | |
| self.down4 = conv4x4(256, 512) | |
| self.down5 = conv4x4(512, 512) | |
| self.up1 = deconv4x4(512, 512) | |
| self.up2 = deconv4x4(512*2, 256) | |
| self.up3 = deconv4x4(256*2, 128) | |
| self.up4 = deconv4x4(128*2, 64) | |
| self.up5 = nn.Conv2d(64*2, 3, kernel_size=3, stride=1, padding=1) | |
| def forward(self, x): # input:(B,6,256,256) | |
| c1 = self.down1(x) | |
| c2 = self.down2(c1) | |
| c3 = self.down3(c2) | |
| c4 = self.down4(c3) | |
| c5 = self.down5(c4) | |
| m1 = self.up1(c5) | |
| m1 = torch.cat((c4, m1), dim=1) | |
| m2 = self.up2(m1) | |
| m2 = torch.cat((c3, m2), dim=1) | |
| m3 = self.up3(m2) | |
| m3 = torch.cat((c2, m3), dim=1) | |
| m4 = self.up4(m3) | |
| m4 = torch.cat((c1, m4), dim=1) | |
| out = nn.functional.interpolate(m4, scale_factor=2, mode='bilinear', align_corners=True) | |
| out = self.up5(out) | |
| return torch.tanh(out) # output:(B,3,256,256) | |
| if __name__ == '__main__': | |
| y_cat = torch.randn(5, 6, 256, 256) | |
| hear = Hear_Net() | |
| y_st = hear(y_cat) | |
| print(y_st.shape) | |