| 145 | } |
| 146 | |
| 147 | torch::Tensor ResNetImpl::features_at(torch::Tensor x, int stage_num) { |
| 148 | assert(stage_num > 0 && "the stage number must in range(1,5)"); |
| 149 | x = conv1->forward(x); |
| 150 | x = bn1->forward(x); |
| 151 | x = torch::relu(x); |
| 152 | if (stage_num == 1) return x; |
| 153 | x = torch::max_pool2d(x, 3, 2, 1); |
| 154 | |
| 155 | x = layer1->forward(x); |
| 156 | if (stage_num == 2) return x; |
| 157 | x = layer2->forward(x); |
| 158 | if (stage_num == 3) return x; |
| 159 | x = layer3->forward(x); |
| 160 | if (stage_num == 4) return x; |
| 161 | x = layer4->forward(x); |
| 162 | if (stage_num == 5) return x; |
| 163 | return x; |
| 164 | } |
| 165 | |
| 166 | void ResNetImpl::load_pretrained(std::string pretrained_path) { |
| 167 | std::map<std::string, std::vector<int>> name2layers = getParams(); |