peachypie98 / rivagan Goto Github PK

Hi, thank you for this amazing work of refactoring the original source code of RivaGAN.

I was wondering if you have any suggestion on how to give more robustness to the model against horizontal flipping attacks.

I trained the model, but after some experiments I noticed that flipping the video leads to very poor accuracy in watermark extraction.

I tried to implement a simple noise layer like this:

class HorizontalFlipping(nn.Module):
    """
    Simulates a transformation where the input is mirrored horizontally.

    Input: (N, 3, L, H, W)
    Output: (N, 3, L, H, W)
    """

    def __init__(self):
        super(HorizontalFlipping, self).__init__()

    def forward(self, frames):
        #return frames.flip(dims=(-1,))
        return frames.flip([4]) #flips along the width dimension

and then I added it inside the training and validation code, similarly to the others noise layers already implemented (Crop, Compression, Scale).

The problem is, when I run a new training with such modifications, the various accuracies have a worsening (included the validation crop accuracy, scale accuracy, ...).

Do you have any suggestions?
Thanks a lot

运行了模型 下载了数据 跑了三百论 loss一直是1.3 精度也不高 只有0.6 左右 这种怎么调节去解决 改变了学习率那个loss也降不下

Hi, thanks for the amazing work.

I was wondering if you could help me in modifying the code so that the model is saved by using its state_dict, instead of saving the whole model, as PyTorch documentation suggests (https://pytorch.org/tutorials/beginner/saving_loading_models.html).

Since the RivaGAN class does not extend directly the nn.Module, I tought to simply modify it from class RivaGAN(object) to class RivaGAN(nn.Module)

and then change from

torch.save(self, os.path.join(log_dir, "model.pt"))

to

torch.save(self.state_dict(), os.path.join(log_dir, "model.pt"))

but maybe I'm missing something and I can't figure out if it's enough for the model to work properly or if this modification could alter the model behavior when loaded (since the RivaGAN class is composed by multiple sub-modules).

Do you have any suggestion?
Thanks

Thank you for sharing the code! Do you have trained checkpoints or is there a way to use their official checkpoint, which ends with .onnx? Thanks!