| 47 | |
| 48 | |
| 49 | std::vector<torch::Tensor> VGGImpl::features(torch::Tensor x, int encoder_depth) { |
| 50 | std::vector<torch::Tensor> ans; |
| 51 | |
| 52 | int j = 0;// layer index of features_ |
| 53 | for (int i = 0; i < cfg.size(); i++) { |
| 54 | if (cfg[i] == -1) { |
| 55 | ans.push_back(x); |
| 56 | if (ans.size() == encoder_depth ) |
| 57 | { |
| 58 | break; |
| 59 | } |
| 60 | x = this->features_[j++]->as<torch::nn::MaxPool2d>()->forward(x); |
| 61 | } |
| 62 | else { |
| 63 | x = this->features_[j++]->as<torch::nn::Conv2d>()->forward(x); |
| 64 | if (batch_norm) { |
| 65 | x = this->features_[j++]->as<torch::nn::BatchNorm2d>()->forward(x); |
| 66 | } |
| 67 | x = this->features_[j++]->as<torch::nn::ReLU>()->forward(x); |
| 68 | } |
| 69 | } |
| 70 | if (ans.size() == encoder_depth && encoder_depth==5) |
| 71 | { |
| 72 | x = this->features_[j++]->as<torch::nn::MaxPool2d>()->forward(x); |
| 73 | ans.push_back(x); |
| 74 | } |
| 75 | return ans; |
| 76 | } |
| 77 | |
| 78 | torch::Tensor VGGImpl::features_at(torch::Tensor x, int stage_num) { |
| 79 | assert(stage_num > 0 && stage_num <=5 && "the stage number must in range[1,5]"); |