Hi,
I am implementing a paper called Hidding Images in Plain Sight: Deep Steganography using three neural networks to hide a secret image in a cover image and then uncover it. My loss function is a combination of two Minimum Squared Errors:
- The difference between the secret image and the output
- The difference between the cover image and the hidden image (with the secret image hidden inside)
However, when training, only the second part of the error is optimized (i.e. trying to make the cover image with the embedded hidden secret image look similar to the original cover image). I have no clue to why this might be happening.
(in order: secret image, output image, cover image and hidden image)
Any clue would be much appreciated!
Here is the NN architectures, loss function and training code.
Neural Network Architecture:
# Preparation Network (2 conv layers)
class PrepNetwork(nn.Module):
def __init__(self):
super(PrepNetwork, self).__init__()
self.layer1P = nn.Sequential(
nn.Conv2d(3, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer2P = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=2),
nn.ReLU(),
nn.Conv2d(50, 3, kernel_size=5, padding=2))
def forward(self, p):
p1 = self.layer1P(p)
out = self.layer2P(p1)
return out
# Hiding Network (5 conv layers)
class HidingNetwork(nn.Module):
def __init__(self):
super(HidingNetwork, self).__init__()
self.layer1H = nn.Sequential(
nn.Conv2d(6, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer2H = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=2),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer3H = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer4H = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=2),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer5H = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
# nn.Conv2d(50, 50, kernel_size=4, padding=2),
# nn.ReLU(),
nn.Conv2d(50, 3, kernel_size=5, padding=2))
def forward(self, h):
h1 = self.layer1H(h)
h2 = self.layer2H(h1)
h3 = self.layer3H(h2)
h4 = self.layer4H(h3)
out = self.layer5H(h4)
return out
# Reveal Network (3 conv layers)
class RevealNetwork(nn.Module):
def __init__(self):
super(RevealNetwork, self).__init__()
self.layer1R = nn.Sequential(
nn.Conv2d(3, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=5, padding=2),
nn.ReLU())
self.layer2R = nn.Sequential(
nn.Conv2d(50, 50, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(50, 50, kernel_size=4, padding=2),
nn.ReLU(),
nn.Conv2d(50, 3, kernel_size=5, padding=2),
nn.ReLU())
self.layer3R = nn.Sequential(
nn.Conv2d(3, 3, kernel_size=1, padding=0))
def forward(self, r):
r1 = self.layer1R(r)
r2 = self.layer2R(r1)
out = self.layer3R(r2)
return out
# Join three networks in one module
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.m1 = PrepNetwork()
self.m2 = HidingNetwork()
self.m3 = RevealNetwork()
def forward(self, secret, cover):
x_1 = self.m1(secret)
mid = torch.cat((x_1, cover), 1)
x_2 = self.m2(mid)
x_3 = self.m3(x_2)
return x_2, x_3
Loss function:
# Implement loss
def customized_loss(S_prime, C_prime, S, C, B):
''' Calculates customized loss.'''
loss_cover = torch.nn.functional.mse_loss(C_prime, C)
loss_secret = torch.nn.functional.mse_loss(S_prime, S)
loss_all = loss_cover + B * loss_secret
return loss_all, loss_cover, loss_secret
Training:
# Save optimizer
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
# Iterate over batches performing forward and backward passes
for epoch in range(num_epochs):
# Train mode
net.train()
# Train one epoch
for idx, train_batch in enumerate(train_loader):
data, _ = train_batch
# Saves secret images and secret covers
train_covers = data[:len(data)//2]
train_secrets = data[len(data)//2:]
# Creates variable from secret and cover images
train_secrets = Variable(train_secrets, requires_grad=False)
train_covers = Variable(train_covers, requires_grad=False)
# Forward + Backward + Optimize
optimizer.zero_grad()
train_hidden, train_output = net(train_secrets, train_covers)
# Calculate loss and perform backprop
train_loss, train_loss_cover, train_loss_secret = customized_loss(train_output, train_hidden, train_secrets, train_covers, beta)
train_loss.backward()
optimizer.step()