| 56 | |
| 57 | |
| 58 | class Net(nn.Module): |
| 59 | def __init__(self, num_classes=751, reid=False): |
| 60 | super(Net, self).__init__() |
| 61 | # 3 128 64 |
| 62 | self.conv = nn.Sequential( |
| 63 | nn.Conv2d(3, 64, 3, stride=1, padding=1), |
| 64 | nn.BatchNorm2d(64), |
| 65 | nn.ReLU(inplace=True), |
| 66 | # nn.Conv2d(32,32,3,stride=1,padding=1), |
| 67 | # nn.BatchNorm2d(32), |
| 68 | # nn.ReLU(inplace=True), |
| 69 | nn.MaxPool2d(3, 2, padding=1), |
| 70 | ) |
| 71 | # 32 64 32 |
| 72 | self.layer1 = make_layers(64, 64, 2, False) |
| 73 | # 32 64 32 |
| 74 | self.layer2 = make_layers(64, 128, 2, True) |
| 75 | # 64 32 16 |
| 76 | self.layer3 = make_layers(128, 256, 2, True) |
| 77 | # 128 16 8 |
| 78 | self.layer4 = make_layers(256, 512, 2, True) |
| 79 | # 256 8 4 |
| 80 | self.avgpool = nn.AvgPool2d((8, 4), 1) |
| 81 | # 256 1 1 |
| 82 | self.reid = reid |
| 83 | self.classifier = nn.Sequential( |
| 84 | nn.Linear(512, 256), |
| 85 | nn.BatchNorm1d(256), |
| 86 | nn.ReLU(inplace=True), |
| 87 | nn.Dropout(), |
| 88 | nn.Linear(256, num_classes), |
| 89 | ) |
| 90 | |
| 91 | def forward(self, x): |
| 92 | x = self.conv(x) |
| 93 | x = self.layer1(x) |
| 94 | x = self.layer2(x) |
| 95 | x = self.layer3(x) |
| 96 | x = self.layer4(x) |
| 97 | x = self.avgpool(x) |
| 98 | x = x.view(x.size(0), -1) |
| 99 | # B x 128 |
| 100 | if self.reid: |
| 101 | x = x.div(x.norm(p=2, dim=1, keepdim=True)) |
| 102 | return x |
| 103 | # classifier |
| 104 | x = self.classifier(x) |
| 105 | return x |
| 106 | |
| 107 | |
| 108 | class Extractor(object): |