| 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) |