| 119 | } |
| 120 | |
| 121 | std::vector<torch::Tensor> ResNetImpl::features(torch::Tensor x, int encoder_depth){ |
| 122 | std::vector<torch::Tensor> features; |
| 123 | features.push_back(x); |
| 124 | x = conv1->forward(x); |
| 125 | x = bn1->forward(x); |
| 126 | x = torch::relu(x); |
| 127 | features.push_back(x); |
| 128 | x = torch::max_pool2d(x, 3, 2, 1); |
| 129 | |
| 130 | std::vector<torch::nn::Sequential> stages = get_stages(); |
| 131 | for (int i = 0; i < encoder_depth - 1; i++) { |
| 132 | x = stages[i]->as<torch::nn::Sequential>()->forward(x); |
| 133 | features.push_back(x); |
| 134 | } |
| 135 | //x = layer1->forward(x); |
| 136 | //features.push_back(x); |
| 137 | //x = layer2->forward(x); |
| 138 | //features.push_back(x); |
| 139 | //x = layer3->forward(x); |
| 140 | //features.push_back(x); |
| 141 | //x = layer4->forward(x); |
| 142 | //features.push_back(x); |
| 143 | |
| 144 | return features; |
| 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)"); |