【发布时间】：2020-05-03 03:30:28
【问题描述】：
我有一个 Tensorflow v1 版本的 unet，它使用 SGD 训练得非常好，学习率为 0.05。
我在 Pytorch 中重写了网络，因为我想公开一些在 Tensorflow 中不那么容易的功能。
我的模型始终预测一个空的蒙版，所以我试图让模型过拟合到一张图像。
可以过拟合一个示例图像来预测一个掩码，但它只适用于 Adam，学习率为 0.0005 和 1000 个 epoch。我的旧模型可以在 10 个 epoch 左右完成。
我看不出任何明显的错误。我一定是做错了什么，因为这是一个微不足道的问题，需要很少的调整。
import numpy as np
import cv2
from PIL import Image
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(42)
class DoubleConv(nn.Module):
"""(convolution => [BN] => ReLU) * 2"""
def __init__(self, in_channels, out_channels, relu=True):
super().__init__()
if relu:
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True)
)
else:
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
)
def forward(self, x):
return self.double_conv(x)
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True, relu=True):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels, relu=relu)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=True):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
self.down4 = Down(512, 512)
self.up1 = Up(1024, 256, bilinear)
self.up2 = Up(512, 128, bilinear)
self.up3 = Up(256, 64, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, 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
def decode_segmap(image, num_classes=3):
label_colors = np.array([(128, 0, 0),
(0, 128, 0), (0, 0, 128)])
r = np.zeros_like(image).astype(np.uint8)
g = np.zeros_like(image).astype(np.uint8)
b = np.zeros_like(image).astype(np.uint8)
for l in range(0, num_classes):
idx = image == l
r[idx] = label_colors[l, 0]
g[idx] = label_colors[l, 1]
b[idx] = label_colors[l, 2]
rgb = np.stack([r, g, b], axis=2)
return rgb
def load_batch(batch_size):
rotated_frame = Image.open('0test.png')
rotated_gt = Image.open('0label.png')
trf = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean = [0.2455], std = [0.2684])])
rotated_frame = trf(rotated_frame).unsqueeze(0)
trf = transforms.Compose([
transforms.ToTensor()])
rotated_gt = trf(rotated_gt).unsqueeze(0)
rotated_frame = torch.mean(rotated_frame, 1).unsqueeze(1)
rotated_gt = torch.mean(rotated_gt, 1).unsqueeze(1)
return rotated_frame.to(device), rotated_gt.type(torch.long).to(device).squeeze(1)
net = UNet(1, 3)
net.to(device=device)
# Loss
#optimizer = optim.RMSprop(net.parameters(), lr=0.005, weight_decay=1e-8)
optimizer = optim.SGD(net.parameters(), lr=0.0005)
#optimizer = optim.Adam(net.parameters(), lr=0.0005)
criterion = nn.CrossEntropyLoss()
# Load data
rotated_frame, rotated_gt = load_batch(1)
print(rotated_frame.shape)
print(rotated_gt.shape)
# Train
epochs = 1000
losses = []
for epoch in range(epochs):
predicted = net(rotated_frame)
loss = criterion(predicted, rotated_gt)
losses.append(loss)
loss.backward()
optimizer.step()
print('Epoch {}/{} Loss: {}'.format(epoch, epochs, loss))
output = torch.argmax(predicted.squeeze(), dim=0).detach().cpu().numpy()
a, b = np.min(output), np.max(output)
print('Predicted: min: {} max: {}'.format(a, b))
print(output.shape)
rgb = decode_segmap(output)
plt.imshow(rgb)
plt.savefig('predicted_argmaxed.png')
gt = rotated_gt.squeeze().detach().cpu().numpy()
a, b = np.min(gt), np.max(gt)
print('Gt: min: {} max: {}'.format(a, b))
rgb = decode_segmap(gt)
plt.imshow(rgb)
plt.savefig('gt_argmaxed.png')
示例图片在这里：
任何帮助将不胜感激！
【问题讨论】：
-
我在你的 DoubleConv 中看不到任何 BN。这可能是罪魁祸首。
-
另外，您确定要在此处使用分类交叉熵吗？如果您想要二元交叉熵，请使用
BCELoss。 -
如果您想在 decode_segmap 输出中看到浮动图像（您的输出示例图像似乎不是二进制的）。如果是这样，请在每个通道中使用 sigmoid 函数。然后进行调试，转储其值并检查图像是否完全为零（或一个）或具有一些浮点值。您可能需要删除 argmax 以得到实际的净输出。
标签： python tensorflow pytorch