class BaseConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding,
                 stride):
        super(BaseConv, self).__init__()

        self.act = nn.ReLU()

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding,
                               stride)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size,
                               padding, stride)

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        return x


class DownConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding,
                 stride):
        super(DownConv, self).__init__()

        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.conv_block = BaseConv(in_channels, out_channels, kernel_size,
                                   padding, stride)

    def forward(self, x):
        x = self.pool1(x)
        x = self.conv_block(x)
        return x


class UpConv(nn.Module):
    def __init__(self, in_channels, in_channels_skip, out_channels,
                 kernel_size, padding, stride):
        super(UpConv, self).__init__()

        self.conv_trans1 = nn.ConvTranspose2d(
            in_channels, in_channels, kernel_size=2, padding=0, stride=2)
        self.conv_block = BaseConv(
            in_channels=in_channels + in_channels_skip,
            out_channels=out_channels,
            kernel_size=kernel_size,
            padding=padding,
            stride=stride)

    def forward(self, x, x_skip):
        x = self.conv_trans1(x)
        x = torch.cat((x, x_skip), dim=1)
        x = self.conv_block(x)
        return x


class UNet(nn.Module):
    def __init__(self, in_channels, out_channels, n_class, kernel_size,
                 padding, stride):
        super(UNet, self).__init__()

        self.init_conv = BaseConv(in_channels, out_channels, kernel_size,
                                  padding, stride)

        self.down1 = DownConv(out_channels, 2 * out_channels, kernel_size,
                              padding, stride)

        self.down2 = DownConv(2 * out_channels, 4 * out_channels, kernel_size,
                              padding, stride)

        self.down3 = DownConv(4 * out_channels, 8 * out_channels, kernel_size,
                              padding, stride)

        self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,
                          kernel_size, padding, stride)

        self.up2 = UpConv(4 * out_channels, 2 * out_channels, 2 * out_channels,
                          kernel_size, padding, stride)

        self.up1 = UpConv(2 * out_channels, out_channels, out_channels,
                          kernel_size, padding, stride)

        self.out = nn.Conv2d(out_channels, n_class, kernel_size, padding, stride)

    def forward(self, x):
        # Encoder
        x = self.init_conv(x)
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        # Decoder
        x_up = self.up3(x3, x2)
        x_up = self.up2(x_up, x1)
        x_up = self.up1(x_up, x)
        x_out = F.sigmoid(self.out(x_up))
        return x_out

model = UNet(in_channels=1,
             out_channels=32,
             n_class=1,
             kernel_size=3,
             padding=1,
             stride=1)
# model = UNet(n_channels=1, n_classes=1)
model = model.to(device)

optimizer = optim.SGD(model.parameters(), lr=0.001)
criterion = nn.BCELoss()

for epoch in range(10):
    model.train()
    for idx, (feat, targ) in enumerate(data.train_loader):
        feat = feat.to(device)
        targ = targ.to(device)

        optimizer.zero_grad()
        pred = model(feat)
        loss = criterion(pred, targ)
        loss.backward()
        optimizer.step()
        
        print(f'Epoch {epoch} Train Loss {idx}: {loss.item()}')
    
    with torch.no_grad():
        model.eval()
        for idx, (feat, targ) in enumerate(data.valid_loader):
            feat = feat.to(device)
            targ = targ.to(device)

            pred = model(feat)
            loss = criterion(pred, targ)

            print(f'Epoch {epoch} Valid Loss {idx}: {loss.item()}')

Epoch 1 Train Loss 37: -2.0091397762298584
Epoch 1 Train Loss 38: -1.8624688386917114
Epoch 1 Train Loss 39: -2.1319470405578613
Epoch 1 Train Loss 40: -2.592771053314209
Epoch 1 Train Loss 41: -3.217898368835449
Epoch 1 Train Loss 42: -3.0293173789978027
Epoch 1 Train Loss 43: -3.380110025405884
Epoch 1 Train Loss 44: -4.6700944900512695
Epoch 1 Valid Loss 0: -8.985795021057129
Epoch 1 Valid Loss 1: -9.246602058410645
Epoch 1 Valid Loss 2: -9.319241523742676
Epoch 1 Valid Loss 3: -9.418668746948242
Epoch 1 Valid Loss 4: -9.03136157989502
Epoch 1 Valid Loss 5: -9.369157791137695
Epoch 1 Valid Loss 6: -9.51211929321289
Epoch 1 Valid Loss 7: -9.425999641418457
Epoch 1 Valid Loss 8: -9.599837303161621
Epoch 1 Valid Loss 9: -9.633952140808105
Epoch 1 Valid Loss 10: -9.373023986816406
Epoch 1 Valid Loss 11: -9.418418884277344
Epoch 1 Valid Loss 12: -9.380928993225098
Epoch 1 Valid Loss 13: -9.855154991149902
Epoch 1 Valid Loss 14: -9.249217987060547