| 145 | } |
| 146 | |
| 147 | void ResNetImpl::make_dilated(std::vector<int> stage_list, std::vector<int> dilation_list) { |
| 148 | if (stage_list.size() != dilation_list.size()) { |
| 149 | std::cout << "make sure stage list len equal to dilation list len"; |
| 150 | return; |
| 151 | } |
| 152 | std::map<int, torch::nn::Sequential> stage_dict = {}; |
| 153 | stage_dict.insert(std::pair<int, torch::nn::Sequential>(5, this->layer4)); |
| 154 | stage_dict.insert(std::pair<int, torch::nn::Sequential>(4, this->layer3)); |
| 155 | stage_dict.insert(std::pair<int, torch::nn::Sequential>(3, this->layer2)); |
| 156 | stage_dict.insert(std::pair<int, torch::nn::Sequential>(2, this->layer1)); |
| 157 | for (int i = 0; i < stage_list.size(); i++) { |
| 158 | int dilation_rate = dilation_list[i]; |
| 159 | for (auto m : stage_dict[stage_list[i]]->modules()) { |
| 160 | if (m->name() == "torch::nn::Conv2dImpl") { |
| 161 | m->as<torch::nn::Conv2d>()->options.stride(1); |
| 162 | m->as<torch::nn::Conv2d>()->options.dilation(dilation_rate); |
| 163 | int kernel_size = m->as<torch::nn::Conv2d>()->options.kernel_size()->at(0); |
| 164 | m->as<torch::nn::Conv2d>()->options.padding((kernel_size / 2) * dilation_rate); |
| 165 | } |
| 166 | } |
| 167 | } |
| 168 | return; |
| 169 | } |
| 170 | |
| 171 | ResNet resnet18(int64_t num_classes) { |
| 172 | std::vector<int> layers = { 2, 2, 2, 2 }; |