| 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]"); |
| 80 | int j = 0; |
| 81 | int stage_count = 0; |
| 82 | for (int i = 0; i < cfg.size(); i++) { |
| 83 | if (cfg[i] == -1) { |
| 84 | x = this->features_[j++]->as<torch::nn::MaxPool2d>()->forward(x); |
| 85 | stage_count++; |
| 86 | if (stage_count == stage_num) |
| 87 | return x; |
| 88 | } |
| 89 | else { |
| 90 | x = this->features_[j++]->as<torch::nn::Conv2d>()->forward(x); |
| 91 | if (batch_norm) { |
| 92 | x = this->features_[j++]->as<torch::nn::BatchNorm2d>()->forward(x); |
| 93 | } |
| 94 | x = this->features_[j++]->as<torch::nn::ReLU>()->forward(x); |
| 95 | } |
| 96 | } |
| 97 | return x; |
| 98 | } |
| 99 | |
| 100 | void VGGImpl::load_pretrained(std::string pretrained_path) { |
| 101 | VGG net_pretrained = VGG(cfg, 1000, batch_norm); |