File size: 2,165 Bytes
d5695ca 49bb7b0 d5695ca 49bb7b0 d5695ca 49bb7b0 d5695ca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import torch
from torch import nn
from transformers import ResNetPreTrainedModel
from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
from transformers.image_processing_utils import BaseImageProcessor
from transformers import ResNetConfig, ResNetModel
from typing import Optional
class ResNetForZeroBitWatermarkDetection(ResNetPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.resnet = ResNetModel(config)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(config.hidden_sizes[-1], 128),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(128, 1))
self.register_buffer('exp', torch.tensor([1.0]))
# initialize weights and apply final processing
self.post_init()
# TODO docstring
def forward(
self,
pixel_values: Optional[torch.FloatTensor] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> ImageClassifierOutputWithNoAttention:
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
outputs = self.resnet(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
pooled_output = outputs.pooler_output if return_dict else outputs[1]
x = self.classifier(pooled_output)
# generalized-Gaussian recalibration, centering and scaling is already included in last linear layer
x = 0.5 + torch.sign(x) * 0.5 * torch.special.gammainc(1 / self.exp, torch.abs(x)**self.exp)
# Laplacian calibration, centering and scaling is already included in last linear layer
# if exp==1
#x = 0.5 + torch.sign(x) * 0.5 * (1 - torch.exp(-torch.abs(x))) # laplacian
logits = torch.log(x) - torch.log1p(-x)
loss = None
if not return_dict:
output = (logits,) + outputs[2:]
return (loss,) + output if loss is not None else output
return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
|