File size: 10,514 Bytes
6d55688 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | import torch
import torch.nn as nn
from torchvision import models
def convrelu(in_channels, out_channels, kernel, padding, pool):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel, padding=padding),
#In conv, the dimension of the output, if the input is H,W, is
# H+2*padding-kernel +1
nn.ReLU(inplace=True),
nn.MaxPool2d(pool, stride=pool, padding=0, dilation=1, return_indices=False, ceil_mode=False)
#pooling takes Height H and width W to (H-pool)/pool+1 = H/pool, and floor. Same for W.
#altogether, the output size is (H+2*padding-kernel +1)/pool.
)
def convreluT(in_channels, out_channels, kernel, padding):
return nn.Sequential(
nn.ConvTranspose2d(in_channels, out_channels, kernel, stride=2, padding=padding),
nn.ReLU(inplace=True)
#input is H X W, output is (H-1)*2 - 2*padding + kernel
)
class RadioWNet(nn.Module):
def __init__(self,inputs=2,phase="firstU"):
super().__init__()
self.inputs=inputs
self.phase=phase
if inputs<=3:
self.layer00 = convrelu(inputs, 6, 3, 1,1)
self.layer0 = convrelu(6, 40, 5, 2,2)
else:
self.layer00 = convrelu(inputs, 10, 3, 1,1)
self.layer0 = convrelu(10, 40, 5, 2,2)
self.layer1 = convrelu(40, 50, 5, 2,2)
self.layer10 = convrelu(50, 60, 5, 2,1)
self.layer2 = convrelu(60, 100, 5, 2,2)
self.layer20 = convrelu(100, 100, 3, 1,1)
self.layer3 = convrelu(100, 150, 5, 2,2)
self.layer4 =convrelu(150, 300, 5, 2,2)
self.layer5 =convrelu(300, 500, 5, 2,2)
self.conv_up5 =convreluT(500, 300, 4, 1)
self.conv_up4 = convreluT(300+300, 150, 4, 1)
self.conv_up3 = convreluT(150 + 150, 100, 4, 1)
self.conv_up20 = convrelu(100 + 100, 100, 3, 1, 1)
self.conv_up2 = convreluT(100 + 100, 60, 6, 2)
self.conv_up10 = convrelu(60 + 60, 50, 5, 2, 1)
self.conv_up1 = convreluT(50 + 50, 40, 6, 2)
self.conv_up0 = convreluT(40 + 40, 20, 6, 2)
if inputs<=3:
self.conv_up00 = convrelu(20+6+inputs, 20, 5, 2,1)
else:
self.conv_up00 = convrelu(20+10+inputs, 20, 5, 2,1)
self.conv_up000 = convrelu(20+inputs, 1, 5, 2,1)
self.Wlayer00 = convrelu(inputs+1, 20, 3, 1,1)
self.Wlayer0 = convrelu(20, 30, 5, 2,2)
self.Wlayer1 = convrelu(30, 40, 5, 2,2)
self.Wlayer10 = convrelu(40, 50, 5, 2,1)
self.Wlayer2 = convrelu(50, 60, 5, 2,2)
self.Wlayer20 = convrelu(60, 70, 3, 1,1)
self.Wlayer3 = convrelu(70, 90, 5, 2,2)
self.Wlayer4 =convrelu(90, 110, 5, 2,2)
self.Wlayer5 =convrelu(110, 150, 5, 2,2)
self.Wconv_up5 =convreluT(150, 110, 4, 1)
self.Wconv_up4 = convreluT(110+110, 90, 4, 1)
self.Wconv_up3 = convreluT(90 + 90, 70, 4, 1)
self.Wconv_up20 = convrelu(70 + 70, 60, 3, 1, 1)
self.Wconv_up2 = convreluT(60 + 60, 50, 6, 2)
self.Wconv_up10 = convrelu(50 + 50, 40, 5, 2, 1)
self.Wconv_up1 = convreluT(40 + 40, 30, 6, 2)
self.Wconv_up0 = convreluT(30 + 30, 20, 6, 2)
self.Wconv_up00 = convrelu(20+20+inputs+1, 20, 5, 2,1)
self.Wconv_up000 = convrelu(20+inputs+1, 1, 5, 2,1)
def forward(self, input):
input0=input[:,0:self.inputs,:,:]
if self.phase=="firstU":
layer00 = self.layer00(input0)
layer0 = self.layer0(layer00)
layer1 = self.layer1(layer0)
layer10 = self.layer10(layer1)
layer2 = self.layer2(layer10)
layer20 = self.layer20(layer2)
layer3 = self.layer3(layer20)
layer4 = self.layer4(layer3)
layer5 = self.layer5(layer4)
layer4u = self.conv_up5(layer5)
layer4u = torch.cat([layer4u, layer4], dim=1)
layer3u = self.conv_up4(layer4u)
layer3u = torch.cat([layer3u, layer3], dim=1)
layer20u = self.conv_up3(layer3u)
layer20u = torch.cat([layer20u, layer20], dim=1)
layer2u = self.conv_up20(layer20u)
layer2u = torch.cat([layer2u, layer2], dim=1)
layer10u = self.conv_up2(layer2u)
layer10u = torch.cat([layer10u, layer10], dim=1)
layer1u = self.conv_up10(layer10u)
layer1u = torch.cat([layer1u, layer1], dim=1)
layer0u = self.conv_up1(layer1u)
layer0u = torch.cat([layer0u, layer0], dim=1)
layer00u = self.conv_up0(layer0u)
layer00u = torch.cat([layer00u, layer00], dim=1)
layer00u = torch.cat([layer00u,input0], dim=1)
layer000u = self.conv_up00(layer00u)
layer000u = torch.cat([layer000u,input0], dim=1)
output1 = self.conv_up000(layer000u)
Winput=torch.cat([output1, input], dim=1).detach()
Wlayer00 = self.Wlayer00(Winput).detach()
Wlayer0 = self.Wlayer0(Wlayer00).detach()
Wlayer1 = self.Wlayer1(Wlayer0).detach()
Wlayer10 = self.Wlayer10(Wlayer1).detach()
Wlayer2 = self.Wlayer2(Wlayer10).detach()
Wlayer20 = self.Wlayer20(Wlayer2).detach()
Wlayer3 = self.Wlayer3(Wlayer20).detach()
Wlayer4 = self.Wlayer4(Wlayer3).detach()
Wlayer5 = self.Wlayer5(Wlayer4).detach()
Wlayer4u = self.Wconv_up5(Wlayer5).detach()
Wlayer4u = torch.cat([Wlayer4u, Wlayer4], dim=1).detach()
Wlayer3u = self.Wconv_up4(Wlayer4u).detach()
Wlayer3u = torch.cat([Wlayer3u, Wlayer3], dim=1).detach()
Wlayer20u = self.Wconv_up3(Wlayer3u).detach()
Wlayer20u = torch.cat([Wlayer20u, Wlayer20], dim=1).detach()
Wlayer2u = self.Wconv_up20(Wlayer20u).detach()
Wlayer2u = torch.cat([Wlayer2u, Wlayer2], dim=1).detach()
Wlayer10u = self.Wconv_up2(Wlayer2u).detach()
Wlayer10u = torch.cat([Wlayer10u, Wlayer10], dim=1).detach()
Wlayer1u = self.Wconv_up10(Wlayer10u).detach()
Wlayer1u = torch.cat([Wlayer1u, Wlayer1], dim=1).detach()
Wlayer0u = self.Wconv_up1(Wlayer1u).detach()
Wlayer0u = torch.cat([Wlayer0u, Wlayer0], dim=1).detach()
Wlayer00u = self.Wconv_up0(Wlayer0u).detach()
Wlayer00u = torch.cat([Wlayer00u, Wlayer00], dim=1).detach()
Wlayer00u = torch.cat([Wlayer00u,Winput], dim=1).detach()
Wlayer000u = self.Wconv_up00(Wlayer00u).detach()
Wlayer000u = torch.cat([Wlayer000u,Winput], dim=1).detach()
output2 = self.Wconv_up000(Wlayer000u).detach()
else:
layer00 = self.layer00(input0).detach()
layer0 = self.layer0(layer00).detach()
layer1 = self.layer1(layer0).detach()
layer10 = self.layer10(layer1).detach()
layer2 = self.layer2(layer10).detach()
layer20 = self.layer20(layer2).detach()
layer3 = self.layer3(layer20).detach()
layer4 = self.layer4(layer3).detach()
layer5 = self.layer5(layer4).detach()
layer4u = self.conv_up5(layer5).detach()
layer4u = torch.cat([layer4u, layer4], dim=1).detach()
layer3u = self.conv_up4(layer4u).detach()
layer3u = torch.cat([layer3u, layer3], dim=1).detach()
layer20u = self.conv_up3(layer3u).detach()
layer20u = torch.cat([layer20u, layer20], dim=1).detach()
layer2u = self.conv_up20(layer20u).detach()
layer2u = torch.cat([layer2u, layer2], dim=1).detach()
layer10u = self.conv_up2(layer2u).detach()
layer10u = torch.cat([layer10u, layer10], dim=1).detach()
layer1u = self.conv_up10(layer10u).detach()
layer1u = torch.cat([layer1u, layer1], dim=1).detach()
layer0u = self.conv_up1(layer1u).detach()
layer0u = torch.cat([layer0u, layer0], dim=1).detach()
layer00u = self.conv_up0(layer0u).detach()
layer00u = torch.cat([layer00u, layer00], dim=1).detach()
layer00u = torch.cat([layer00u,input0], dim=1).detach()
layer000u = self.conv_up00(layer00u).detach()
layer000u = torch.cat([layer000u,input0], dim=1).detach()
output1 = self.conv_up000(layer000u).detach()
Winput=torch.cat([output1, input], dim=1).detach()
Wlayer00 = self.Wlayer00(Winput)
Wlayer0 = self.Wlayer0(Wlayer00)
Wlayer1 = self.Wlayer1(Wlayer0)
Wlayer10 = self.Wlayer10(Wlayer1)
Wlayer2 = self.Wlayer2(Wlayer10)
Wlayer20 = self.Wlayer20(Wlayer2)
Wlayer3 = self.Wlayer3(Wlayer20)
Wlayer4 = self.Wlayer4(Wlayer3)
Wlayer5 = self.Wlayer5(Wlayer4)
Wlayer4u = self.Wconv_up5(Wlayer5)
Wlayer4u = torch.cat([Wlayer4u, Wlayer4], dim=1)
Wlayer3u = self.Wconv_up4(Wlayer4u)
Wlayer3u = torch.cat([Wlayer3u, Wlayer3], dim=1)
Wlayer20u = self.Wconv_up3(Wlayer3u)
Wlayer20u = torch.cat([Wlayer20u, Wlayer20], dim=1)
Wlayer2u = self.Wconv_up20(Wlayer20u)
Wlayer2u = torch.cat([Wlayer2u, Wlayer2], dim=1)
Wlayer10u = self.Wconv_up2(Wlayer2u)
Wlayer10u = torch.cat([Wlayer10u, Wlayer10], dim=1)
Wlayer1u = self.Wconv_up10(Wlayer10u)
Wlayer1u = torch.cat([Wlayer1u, Wlayer1], dim=1)
Wlayer0u = self.Wconv_up1(Wlayer1u)
Wlayer0u = torch.cat([Wlayer0u, Wlayer0], dim=1)
Wlayer00u = self.Wconv_up0(Wlayer0u)
Wlayer00u = torch.cat([Wlayer00u, Wlayer00], dim=1)
Wlayer00u = torch.cat([Wlayer00u,Winput], dim=1)
Wlayer000u = self.Wconv_up00(Wlayer00u)
Wlayer000u = torch.cat([Wlayer000u,Winput], dim=1)
output2 = self.Wconv_up000(Wlayer000u)
return [output1,output2]
|