From cdb00ebd8dfc3f2f05193422927173747adcf409 Mon Sep 17 00:00:00 2001 From: Yin Li Date: Wed, 22 Jan 2020 14:01:18 -0500 Subject: [PATCH] Add a simple PatchGAN --- map2map/models/patchgan.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 map2map/models/patchgan.py diff --git a/map2map/models/patchgan.py b/map2map/models/patchgan.py new file mode 100644 index 0000000..74a59af --- /dev/null +++ b/map2map/models/patchgan.py @@ -0,0 +1,18 @@ +import torch.nn as nn + +from .conv import ConvBlock + + +class PatchGAN(nn.Module): + def __init__(self, in_chan, out_chan=1): + super().__init__() + + self.convs = nn.Sequential( + ConvBlock(in_chan, 64, seq='CA'), + ConvBlock(64, 128, seq='CBA'), + ConvBlock(128, 256, seq='CBA'), + nn.Conv3d(256, out_chan, 1) + ) + + def forward(self, x): + return self.convs(x)