Adjust BitNetwork Message Expanding

[huayuan4396]

I met a problem where the bit error won't decrease when training on 384 * 384 cover image with 128 bits of message length.

I found the problem is aroused since the message is first expanded to [b, 1, msg_len, msg_len] and than interpolated to the feature map size [b, 128, img_size // 8, img_size // 8]. In this procedure, the expanded message tensor is downsampled and can cause information loss when msg_len is larger than img_size // 8 (in my case, msg_len = 128 and img_size // 8 = 48).

As a result, I tried some modification and found the submitted changes can solve this problem, in which I choose to directly expand the message tensor to [b, 1, img_size // 8, img_size // 8] instead of [b, 1, msg_len, msg_len]. The solution is simple and effective, by just changing one linear layer. The bit error rate is reduced quickly within 5 epoches to 0.0273:

Hope this is can help improve the model:)