Coverage for robustAI/advertrain/models.py: 42%

1import torch 

2from torch import Tensor, nn 

3from torch.nn import functional as F 

4 

5from robustAI.advertrain.dependencies.dropblock import DropBlock2d 

6 

7 

8class Normalize(nn.Module): 

9 def __init__(self, mean: Tensor, std: Tensor, device: torch.device) -> None: 

10 """ 

11 Initialize the Normalize module. 

12 

13 This module is used to normalize image data by subtracting the mean and 

14 dividing by the standard deviation. 

15 

16 Args: 

17 mean (Tensor): A tensor containing the mean values for each channel. 

18 std (Tensor): A tensor containing the standard deviation for each channel. 

19 device (torch.device): The device (CPU or GPU) to which the tensors should be moved. 

20 """ 

21 super().__init__() 

22 

23 self.mean = mean.unsqueeze(-1).unsqueeze(-1).to(device) 

24 self.std = std.unsqueeze(-1).unsqueeze(-1).to(device) 

25 

26 def forward(self, x: Tensor) -> Tensor: 

27 """ 

28 Normalize the input tensor. 

29 

30 Applies the normalization operation on the input tensor using the mean and 

31 standard deviation provided during initialization. 

32 

33 Args: 

34 x (Tensor): The input tensor to be normalized. 

35 

36 Returns: 

37 Tensor: The normalized tensor. 

38 """ 

39 return (x - self.mean) / self.std 

40 

41 

42class ConvNet(nn.Module): 

43 """ 

44 Convolutional Neural Network with dropout layers, designed for processing images of size 64x128. 

45 

46 This network includes a normalization layer, several convolutional layers with 

47 ReLU activation and max pooling, followed by fully connected layers with dropout 

48 for regularization. It is suited for tasks like image classification where dropout 

49 can help reduce overfitting. 

50 

51 Attributes: 

52 norm (Normalize): Normalization layer to preprocess the input images. 

53 conv1, conv2_1, conv3_1, conv4_1 (nn.Conv2d): Convolutional layers for feature extraction. 

54 pooling (nn.MaxPool2d): Max pooling layer to reduce spatial dimensions. 

55 activation (nn.ReLU): Activation function. 

56 dropout (nn.Dropout): Dropout layer for regularization. 

57 linear1, linear2, linear3 (nn.Linear): Fully connected layers for classification. 

58 """ 

59 

60 def __init__(self, device: torch.device, p: float = 0.2) -> None: 

61 """ 

62 Initializes the ConvNetDropout model with dropout layers. 

63 

64 Args: 

65 device (torch.device): The device to which the model and tensors should be moved. 

66 p (float): Dropout probability. Default is 0.2. 

67 """ 

68 super().__init__() 

69 

70 self.norm = Normalize( 

71 torch.Tensor([0.4632, 0.4532, 0.4485]), 

72 torch.Tensor([0.1646, 0.1759, 0.1739]), 

73 device, 

74 ) 

75 

76 self.conv1 = nn.Conv2d(3, 32, 7, padding=3) 

77 self.conv2_1 = nn.Conv2d(32, 64, 5, padding=2) 

78 self.conv3_1 = nn.Conv2d(64, 128, 5, padding=2) 

79 self.conv4_1 = nn.Conv2d(128, 256, 5, padding=2) 

80 

81 self.pooling = nn.MaxPool2d(2) 

82 self.activation = nn.ReLU() 

83 self.dropout = nn.Dropout(p=p) 

84 

85 self.linear1 = nn.Linear(4 * 8 * 256, 2048) 

86 self.linear2 = nn.Linear(2048, 1024) 

87 self.linear3 = nn.Linear(1024, 2) 

88 

89 self.to(device) 

90 

91 def forward(self, x: Tensor) -> Tensor: 

92 """ 

93 Defines the forward pass of the ConvNetDropout. 

94 

95 The input tensor is processed through normalization, convolutional layers, 

96 pooling layers, dropout layers, and fully connected layers sequentially to 

97 produce the output tensor. 

98 

99 Args: 

100 x (Tensor): Input tensor of shape (batch_size, 3, 64, 128). 

101 

102 Returns: 

103 Tensor: Output tensor after processing through the network. 

104 """ 

105 x = self.norm(x) 

106 

107 y = self.activation(self.conv1(x)) 

108 y = self.pooling(y) 

109 

110 y = self.activation(self.conv2_1(y)) 

111 y = self.pooling(y) 

112 

113 y = self.activation(self.conv3_1(y)) 

114 y = self.pooling(y) 

115 

116 y = self.activation(self.conv4_1(y)) 

117 y = self.pooling(y) 

118 

119 y = self.activation(self.linear1(torch.reshape(y, (-1, 4 * 8 * 256)))) 

120 y = self.dropout(y) 

121 y = self.activation(self.linear2(y)) 

122 y = self.dropout(y) 

123 y = self.linear3(y) 

124 

125 return y 

126 

127 

128class ConvNetDropblock(nn.Module): 

129 """ 

130 Convolutional Neural Network with DropBlock regularization, designed for processing images of size 64x128. 

131 

132 This network includes a normalization layer, several convolutional layers with 

133 ReLU activation and max pooling, followed by fully connected layers with dropout 

134 and DropBlock for regularization. It is suited for tasks like image classification 

135 where advanced regularization techniques can help reduce overfitting. 

136 

137 Attributes: 

138 norm (Normalize): Normalization layer to preprocess the input images. 

139 conv1, conv2_1, conv3_1, conv4_1 (nn.Conv2d): Convolutional layers for feature extraction. 

140 pooling (nn.MaxPool2d): Max pooling layer to reduce spatial dimensions. 

141 activation (nn.ReLU): Activation function. 

142 dropout (nn.Dropout): Dropout layer for regularization. 

143 dropblock (DropBlock2d): DropBlock layer for structured dropout. 

144 linear1, linear2, linear3 (nn.Linear): Fully connected layers for classification. 

145 """ 

146 

147 def __init__(self, device: torch.device, p: float = 0.2, drop_prob: float = 0.0, n_steps: int = 10) -> None: 

148 """ 

149 Initializes the ConvNetDropblock model with DropBlock layers. 

150 

151 Args: 

152 device (torch.device): The device (CPU or GPU) to which the model and tensors should be moved. 

153 p (float): Dropout probability for the standard dropout layers. Default is 0.2. 

154 drop_prob (float): Initial probability for DropBlock. Default is 0.1. 

155 n_steps (int): Number of steps over which DropBlock probability should reach its maximum. Default is 10. 

156 """ 

157 super().__init__() 

158 

159 self.norm = Normalize( 

160 torch.Tensor([0.4632, 0.4532, 0.4485]), 

161 torch.Tensor([0.1646, 0.1759, 0.1739]), 

162 device, 

163 ) 

164 

165 self.conv1 = nn.Conv2d(3, 32, 7, padding=3) 

166 self.conv2_1 = nn.Conv2d(32, 64, 5, padding=2) 

167 self.conv3_1 = nn.Conv2d(64, 128, 5, padding=2) 

168 self.conv4_1 = nn.Conv2d(128, 256, 5, padding=2) 

169 

170 self.pooling = nn.MaxPool2d(2) 

171 self.activation = nn.ReLU() 

172 self.dropout = nn.Dropout(p=p) 

173 

174 self.dropblock = DropBlock2d(drop_prob=drop_prob) 

175 self.drop_prob = drop_prob 

176 self.n_epochs = n_steps 

177 self.epochs = 0 

178 

179 self.linear1 = nn.Linear(4 * 8 * 256, 2048) 

180 self.linear2 = nn.Linear(2048, 1024) 

181 self.linear3 = nn.Linear(1024, 2) 

182 

183 self.to(device) 

184 

185 def forward(self, x: Tensor) -> Tensor: 

186 """ 

187 Defines the forward pass of the ConvNetDropblock. 

188 

189 The input tensor is processed through normalization, convolutional layers, 

190 pooling layers, DropBlock layers, dropout layers, and fully connected layers 

191 sequentially to produce the output tensor. 

192 

193 Args: 

194 x (Tensor): Input tensor of shape. 

195 

196 Returns: 

197 Tensor: Output tensor after processing through the network. 

198 """ 

199 x = self.norm(x) 

200 

201 y = self.activation(self.conv1(x)) 

202 y = self.pooling(y) 

203 y = self.dropblock(y) 

204 

205 y = self.activation(self.conv2_1(y)) 

206 y = self.pooling(y) 

207 y = self.dropblock(y) 

208 

209 y = self.activation(self.conv3_1(y)) 

210 y = self.pooling(y) 

211 y = self.dropblock(y) 

212 

213 y = self.activation(self.conv4_1(y)) 

214 y = self.pooling(y) 

215 

216 y = self.activation(self.linear1(torch.reshape(y, (-1, 4 * 8 * 256)))) 

217 y = self.dropout(y) 

218 y = self.activation(self.linear2(y)) 

219 y = self.dropout(y) 

220 y = self.linear3(y) 

221 

222 return y 

223 

224 

225class ResNet(nn.Module): 

226 """ 

227 A custom implementation of a Residual Network (ResNet) for processing images. 

228 

229 This network consists of multiple convolutional layers, each followed by batch normalization, 

230 and some layers include dropout for regularization. The network uses skip connections 

231 similar to a ResNet architecture, adding the output of one layer to another layer. 

232 """ 

233 

234 def __init__(self, device: torch.device, p: float = 0.2) -> None: 

235 """ 

236 Initializes the ResNet model. 

237 

238 Args: 

239 device (torch.device): The device to which the model and tensors should be moved. 

240 """ 

241 super().__init__() 

242 

243 self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) 

244 self.conv1_bn = nn.BatchNorm2d(32) 

245 self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=1) 

246 self.conv2_bn = nn.BatchNorm2d(32) 

247 self.conv3 = nn.Conv2d(32, 32, kernel_size=3, padding=1) 

248 self.conv3_drop = nn.Dropout2d(p=0.2) 

249 self.conv3_bn = nn.BatchNorm2d(32) 

250 

251 self.conv4 = nn.Conv2d(32, 64, kernel_size=3, padding=1) 

252 self.conv4_bn = nn.BatchNorm2d(64) 

253 self.conv5 = nn.Conv2d(64, 64, kernel_size=3, padding=1) 

254 self.conv5_bn = nn.BatchNorm2d(64) 

255 self.conv6 = nn.Conv2d(64, 64, kernel_size=3, padding=1) 

256 self.conv6_drop = nn.Dropout2d(p=0.2) 

257 self.conv6_bn = nn.BatchNorm2d(64) 

258 

259 self.conv7 = nn.Conv2d(64, 128, kernel_size=3, padding=1) 

260 self.conv7_bn = nn.BatchNorm2d(128) 

261 self.conv8 = nn.Conv2d(128, 128, kernel_size=3, padding=1) 

262 self.conv8_bn = nn.BatchNorm2d(128) 

263 self.conv9 = nn.Conv2d(128, 128, kernel_size=3, padding=1) 

264 self.conv9_drop = nn.Dropout2d(p=p) 

265 self.conv9_bn = nn.BatchNorm2d(128) 

266 

267 self.conv10 = nn.Conv2d(128, 256, kernel_size=3, padding=1) 

268 self.conv10_bn = nn.BatchNorm2d(256) 

269 self.conv11 = nn.Conv2d(256, 256, kernel_size=3, padding=1) 

270 self.conv11_bn = nn.BatchNorm2d(256) 

271 self.conv12 = nn.Conv2d(256, 256, kernel_size=3, padding=1) 

272 self.conv12_drop = nn.Dropout2d(p=p) 

273 self.conv12_bn = nn.BatchNorm2d(256) 

274 

275 self.fc1 = nn.Linear(256 * 8 * 16, 2048) 

276 self.fc1_bn = nn.BatchNorm1d(2048) 

277 self.fc2 = nn.Linear(2048, 2) 

278 

279 self.to(device) 

280 

281 def forward(self, inp: Tensor) -> Tensor: 

282 """ 

283 Defines the forward pass of the ResNet. 

284 

285 The input tensor is processed through a series of convolutional layers with skip connections, 

286 batch normalization, and dropout, followed by fully connected layers to produce the output tensor. 

287 

288 Args: 

289 inp (Tensor): Input tensor of appropriate shape, typically matching the input size of the first 

290 convolutional layer. 

291 

292 Returns: 

293 Tensor: Output tensor after processing through the network. 

294 """ 

295 res = F.relu(self.conv1_bn(self.conv1(inp))) 

296 x = F.relu(self.conv2_bn(self.conv2(res))) 

297 x = self.conv3_drop(self.conv3(x)) 

298 block1_out = F.relu(self.conv3_bn(F.max_pool2d(x + res, 2))) # 32x64 

299 

300 res = F.relu(self.conv4_bn(self.conv4(block1_out))) 

301 x = F.relu(self.conv5_bn(self.conv5(res))) 

302 x = self.conv6_drop(self.conv6(x)) 

303 block2_out = F.relu(self.conv6_bn(F.max_pool2d(x + res, 2))) # 16x32 

304 

305 res = F.relu(self.conv7_bn(self.conv7(block2_out))) 

306 x = F.relu(self.conv8_bn(self.conv8(res))) 

307 x = self.conv9_drop(self.conv9(x)) 

308 block3_out = F.relu(self.conv9_bn(F.max_pool2d(x + res, 2))) # 8x16 

309 

310 res = F.relu(self.conv10_bn(self.conv10(block3_out))) 

311 x = F.relu(self.conv11_bn(self.conv11(res))) 

312 x = F.relu(self.conv12_bn(self.conv12_drop(self.conv12(x + res)))) 

313 

314 x = x.view(-1, 256 * 8 * 16) 

315 x = F.relu(self.fc1_bn(self.fc1(x))) 

316 x = F.dropout(x, training=self.training, p=0.2) 

317 x = self.fc2(x) 

318 return x 

319 

320 

321class ResNetDropblock(nn.Module): 

322 """ 

323 A custom implementation of a Residual Network (ResNet) for processing images. 

324 

325 This network consists of multiple convolutional layers, each followed by batch normalization, 

326 and some layers include dropout for regularization. The network uses skip connections 

327 similar to a ResNet architecture, adding the output of one layer to another layer. 

328 """ 

329 

330 def __init__(self, device: torch.device, p: float = 0.2, drop_prob: float = 0.0) -> None: 

331 """ 

332 Initializes the ResNet model. 

333 

334 Args: 

335 device (torch.device): The device to which the model and tensors should be moved. 

336 """ 

337 super().__init__() 

338 

339 self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) 

340 self.conv1_bn = nn.BatchNorm2d(32) 

341 self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=1) 

342 self.conv2_bn = nn.BatchNorm2d(32) 

343 self.conv3 = nn.Conv2d(32, 32, kernel_size=3, padding=1) 

344 self.conv3_drop = nn.Dropout2d(p=0.2) 

345 self.conv3_bn = nn.BatchNorm2d(32) 

346 

347 self.conv4 = nn.Conv2d(32, 64, kernel_size=3, padding=1) 

348 self.conv4_bn = nn.BatchNorm2d(64) 

349 self.conv5 = nn.Conv2d(64, 64, kernel_size=3, padding=1) 

350 self.conv5_bn = nn.BatchNorm2d(64) 

351 self.conv6 = nn.Conv2d(64, 64, kernel_size=3, padding=1) 

352 self.conv6_drop = nn.Dropout2d(p=0.2) 

353 self.conv6_bn = nn.BatchNorm2d(64) 

354 

355 self.conv7 = nn.Conv2d(64, 128, kernel_size=3, padding=1) 

356 self.conv7_bn = nn.BatchNorm2d(128) 

357 self.conv8 = nn.Conv2d(128, 128, kernel_size=3, padding=1) 

358 self.conv8_bn = nn.BatchNorm2d(128) 

359 self.conv9 = nn.Conv2d(128, 128, kernel_size=3, padding=1) 

360 self.conv9_drop = nn.Dropout2d(p=p) 

361 self.conv9_bn = nn.BatchNorm2d(128) 

362 

363 self.conv10 = nn.Conv2d(128, 256, kernel_size=3, padding=1) 

364 self.conv10_bn = nn.BatchNorm2d(256) 

365 self.conv11 = nn.Conv2d(256, 256, kernel_size=3, padding=1) 

366 self.conv11_bn = nn.BatchNorm2d(256) 

367 self.conv12 = nn.Conv2d(256, 256, kernel_size=3, padding=1) 

368 self.conv12_drop = nn.Dropout2d(p=p) 

369 self.conv12_bn = nn.BatchNorm2d(256) 

370 

371 self.dropblock = DropBlock2d(drop_prob=drop_prob) 

372 self.drop_prob = drop_prob 

373 

374 self.fc1 = nn.Linear(256 * 8 * 16, 2048) 

375 self.fc1_bn = nn.BatchNorm1d(2048) 

376 self.fc2 = nn.Linear(2048, 2) 

377 

378 self.to(device) 

379 

380 def forward(self, inp: Tensor) -> Tensor: 

381 """ 

382 Defines the forward pass of the ResNet. 

383 

384 The input tensor is processed through a series of convolutional layers with skip connections, 

385 batch normalization, and dropout, followed by fully connected layers to produce the output tensor. 

386 

387 Args: 

388 inp (Tensor): Input tensor of appropriate shape, typically matching the input size of the first 

389 convolutional layer. 

390 

391 Returns: 

392 Tensor: Output tensor after processing through the network. 

393 """ 

394 res = F.relu(self.conv1_bn(self.conv1(inp))) 

395 x = F.relu(self.conv2_bn(self.conv2(res))) 

396 x = self.conv3_drop(self.conv3(x)) 

397 x = self.dropblock(x) 

398 block1_out = F.relu(self.conv3_bn(F.max_pool2d(x + res, 2))) # 32x64 

399 

400 res = F.relu(self.conv4_bn(self.conv4(block1_out))) 

401 x = F.relu(self.conv5_bn(self.conv5(res))) 

402 x = self.conv6_drop(self.conv6(x)) 

403 x = self.dropblock(x) 

404 block2_out = F.relu(self.conv6_bn(F.max_pool2d(x + res, 2))) # 16x32 

405 

406 res = F.relu(self.conv7_bn(self.conv7(block2_out))) 

407 x = F.relu(self.conv8_bn(self.conv8(res))) 

408 x = self.conv9_drop(self.conv9(x)) 

409 x = self.dropblock(x) 

410 block3_out = F.relu(self.conv9_bn(F.max_pool2d(x + res, 2))) # 8x16 

411 

412 res = F.relu(self.conv10_bn(self.conv10(block3_out))) 

413 x = F.relu(self.conv11_bn(self.conv11(res))) 

414 x = F.relu(self.conv12_bn(self.conv12_drop(self.conv12(x + res)))) 

415 x = self.dropblock(x) 

416 

417 x = x.view(-1, 256 * 8 * 16) 

418 x = F.relu(self.fc1_bn(self.fc1(x))) 

419 x = F.dropout(x, training=self.training, p=0.2) 

420 x = self.fc2(x) 

421 return x