Laurent/REVA-QCAV
1
0
Fork 0
mirror of https://github.com/Laurent2916/REVA-QCAV.git synced 2024-09-19 11:35:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: made model.py extensible

Browse source 
Former-commit-id: 0261a3f6caf571dc19e6ca97d2ecada1e72c7f04
This commit is contained in:
Your Name 2022-06-27 16:31:07 +02:00
parent bc2aeacfa3
commit a42190ec61
1 changed files with 16 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
src/unet/model.py
View file

@ -20,20 +20,24 @@ class UNet(nn.Module):
self.ups = nn.ModuleList()
for i in range(len(features) - 1):
self.ups.append(
Up(*features[::-1][i : i + 2]),
Up(*features[-1 - i : -1 - i + 3 : -1]),
)
self.outc = OutConv(features[0], n_classes)
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = self.outc(x)
return logits
skips = []
x = self.inc(x)
for down in self.downs:
skips.append(x)
x = down(x)
for up, skip in zip(self.ups, reversed(skips)):
x = up(x, skip)
x = self.outc(x)
return x