From 31ceb97996b655451bcf26046e29dbbb895d6881 Mon Sep 17 00:00:00 2001 From: Laurent Fainsin Date: Fri, 1 Jul 2022 10:33:35 +0200 Subject: [PATCH] fix: removed now useless Sequential Former-commit-id: 4b97d6d3bef3814f7cfc6c6ceca7833105b4211a [formerly 29f131cf14a73b22427f4ef2318def6870cbe7e0] Former-commit-id: db11e7b2ef2ba39fe11e69251731ea2489d3be31 --- src/train.py | 4 ++-- src/unet/blocks.py | 5 +---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/train.py b/src/train.py index de2fd48..4a434ed 100644 --- a/src/train.py +++ b/src/train.py @@ -22,7 +22,7 @@ def main(): wandb.init( project="U-Net", config=dict( - DIR_TRAIN_IMG="/home/lilian/data_disk/lfainsin/smolval2017", + DIR_TRAIN_IMG="/home/lilian/data_disk/lfainsin/val2017", DIR_VALID_IMG="/home/lilian/data_disk/lfainsin/smoltrain2017/", DIR_SPHERE_IMG="/home/lilian/data_disk/lfainsin/spheres/Images/", DIR_SPHERE_MASK="/home/lilian/data_disk/lfainsin/spheres/Masks/", @@ -31,7 +31,7 @@ def main(): N_CLASSES=1, AMP=True, PIN_MEMORY=True, - BENCHMARK=False, + BENCHMARK=True, DEVICE="cuda", WORKERS=8, EPOCHS=5, diff --git a/src/unet/blocks.py b/src/unet/blocks.py index b1f7d7f..1f4a854 100644 --- a/src/unet/blocks.py +++ b/src/unet/blocks.py @@ -70,10 +70,7 @@ class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super(OutConv, self).__init__() - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=1), - # nn.Sigmoid(), - ) + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1) def forward(self, x): return self.conv(x)