Merge pull request #37 from rht/master

Move the sigmoid activation to the model itself

Former-commit-id: 7dd7c8b6346033ed1b15a2c1ab54847016a826db
This commit is contained in:
milesial 2018-11-11 20:56:15 +01:00 committed by GitHub
commit 5da5e17c08
4 changed files with 7 additions and 7 deletions

View file

@ -20,7 +20,7 @@ def eval_net(net, dataset, gpu=False):
true_mask = true_mask.cuda()
mask_pred = net(img)[0]
mask_pred = (F.sigmoid(mask_pred) > 0.5).float()
mask_pred = (mask_pred > 0.5).float()
tot += dice_coeff(mask_pred, true_mask).item()
return tot / i

View file

@ -43,8 +43,8 @@ def predict_img(net,
output_left = net(X_left)
output_right = net(X_right)
left_probs = F.sigmoid(output_left).squeeze(0)
right_probs = F.sigmoid(output_right).squeeze(0)
left_probs = output_left.squeeze(0)
right_probs = output_right.squeeze(0)
tf = transforms.Compose(
[

View file

@ -6,7 +6,6 @@ import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from eval import eval_net
@ -74,8 +73,7 @@ def train_net(net,
true_masks = true_masks.cuda()
masks_pred = net(imgs)
masks_probs = F.sigmoid(masks_pred)
masks_probs_flat = masks_probs.view(-1)
masks_probs_flat = masks_pred.view(-1)
true_masks_flat = true_masks.view(-1)

View file

@ -1,5 +1,7 @@
# full assembly of the sub-parts to form the complete net
import torch.nn.functional as F
from .unet_parts import *
class UNet(nn.Module):
@ -27,4 +29,4 @@ class UNet(nn.Module):
x = self.up3(x, x2)
x = self.up4(x, x1)
x = self.outc(x)
return x
return F.sigmoid(x)