(self, is_pretrain = True, size1 = 512, size2 = 256)
| 78 | |
| 79 | class ICNet(nn.Module): |
| 80 | def __init__(self, is_pretrain = True, size1 = 512, size2 = 256): |
| 81 | super(ICNet,self).__init__() |
| 82 | resnet18Pretrained1 = torchvision.models.resnet18(pretrained= is_pretrain) |
| 83 | resnet18Pretrained2 = torchvision.models.resnet18(pretrained= is_pretrain) |
| 84 | |
| 85 | self.size1 = size1 |
| 86 | self.size2 = size2 |
| 87 | |
| 88 | ## detail branch |
| 89 | self.b1_1 = nn.Sequential(*list(resnet18Pretrained1.children())[:5]) |
| 90 | self.b1_1_slam = slam(32) |
| 91 | |
| 92 | self.b1_2 = list(resnet18Pretrained1.children())[5] |
| 93 | self.b1_2_slam = slam(32) |
| 94 | |
| 95 | ## context branch |
| 96 | self.b2_1 = nn.Sequential(*list(resnet18Pretrained2.children())[:5]) |
| 97 | self.b2_1_slam = slam(32) |
| 98 | |
| 99 | self.b2_2 = list(resnet18Pretrained2.children())[5] |
| 100 | self.b2_2_slam = slam(32) |
| 101 | |
| 102 | self.b2_3 = list(resnet18Pretrained2.children())[6] |
| 103 | self.b2_3_slam = slam(16) |
| 104 | |
| 105 | self.b2_4 = list(resnet18Pretrained2.children())[7] |
| 106 | self.b2_4_slam = slam(8) |
| 107 | |
| 108 | ## upsample |
| 109 | self.upsize = size1 // 8 |
| 110 | self.up1 = up_conv_bn_relu(up_size = self.upsize, in_channels = 128, out_channels = 256) |
| 111 | self.up2 = up_conv_bn_relu(up_size = self.upsize, in_channels = 512, out_channels = 256) |
| 112 | |
| 113 | ## map prediction head |
| 114 | self.to_map_f = conv_bn_relu(256*2,256*2) |
| 115 | self.to_map_f_slam = slam(32) |
| 116 | self.to_map = to_map(256*2) |
| 117 | |
| 118 | ## score prediction head |
| 119 | self.to_score_f = conv_bn_relu(256*2,256*2) |
| 120 | self.to_score_f_slam = slam(32) |
| 121 | self.head = nn.Sequential( |
| 122 | nn.Linear(256*2,512), |
| 123 | nn.ReLU(), |
| 124 | nn.Linear(512,1), |
| 125 | nn.Sigmoid() |
| 126 | ) |
| 127 | self.avgpool = nn.AdaptiveAvgPool2d((1,1)) |
| 128 | |
| 129 | |
| 130 | def forward(self,x1): |
nothing calls this directly
no test coverage detected