| 114 | |
| 115 | class WhiteBoxAttackModel(nn.Module): |
| 116 | def __init__(self, class_num, embedding_dim): |
| 117 | super(WhiteBoxAttackModel, self).__init__() |
| 118 | |
| 119 | self.dropout = nn.Dropout(p=0.2) |
| 120 | self.output_component = nn.Sequential( |
| 121 | nn.Linear(class_num, 128), |
| 122 | nn.ReLU(), |
| 123 | nn.Linear(128, 64), |
| 124 | ) |
| 125 | |
| 126 | self.loss_component = nn.Sequential( |
| 127 | nn.Linear(1, 128), |
| 128 | nn.ReLU(), |
| 129 | nn.Linear(128, 64), |
| 130 | ) |
| 131 | |
| 132 | self.gradient_component = nn.Sequential( |
| 133 | nn.Conv2d(1, 5, kernel_size=3, padding=0), |
| 134 | nn.AdaptiveAvgPool2d(5), # [batch_size, channel, 5, 5] |
| 135 | nn.Flatten(), |
| 136 | nn.Linear(125, 256), |
| 137 | nn.ReLU(), |
| 138 | nn.Linear(256, 128), |
| 139 | nn.ReLU(), |
| 140 | nn.Linear(128, 64), |
| 141 | ) |
| 142 | |
| 143 | self.label_component = nn.Sequential( |
| 144 | nn.Linear(class_num, 128), |
| 145 | nn.ReLU(), |
| 146 | nn.Linear(128, 64), |
| 147 | ) |
| 148 | |
| 149 | self.embedding_component = nn.Sequential( |
| 150 | nn.Linear(embedding_dim, 128), |
| 151 | nn.ReLU(), |
| 152 | nn.Linear(128, 64), |
| 153 | ) |
| 154 | |
| 155 | self.encoder_component = nn.Sequential( |
| 156 | nn.Linear(64 * 4, 256), |
| 157 | nn.ReLU(), |
| 158 | nn.Linear(256, 128), |
| 159 | nn.ReLU(), |
| 160 | nn.Linear(128, 64), |
| 161 | nn.ReLU(), |
| 162 | nn.Linear(64, 2), |
| 163 | ) |
| 164 | |
| 165 | def forward(self, original_label, output, gradient, embedding, loss): |
| 166 | label_component_result = self.label_component(original_label) |