| 1 | #include "vgg.h" |
| 2 | |
| 3 | torch::nn::Sequential make_features(std::vector<int> &cfg, bool batch_norm){ |
| 4 | torch::nn::Sequential features; |
| 5 | int in_channels = 3; |
| 6 | for(auto v : cfg){ |
| 7 | if(v==-1){ |
| 8 | features->push_back(torch::nn::MaxPool2d(maxpool_options(2,2))); |
| 9 | } |
| 10 | else{ |
| 11 | auto conv2d = torch::nn::Conv2d(conv_options(in_channels,v,3,1,1)); |
| 12 | features->push_back(conv2d); |
| 13 | if(batch_norm){ |
| 14 | features->push_back(torch::nn::BatchNorm2d(torch::nn::BatchNorm2dOptions(v))); |
| 15 | } |
| 16 | features->push_back(torch::nn::ReLU(torch::nn::ReLUOptions(true))); |
| 17 | in_channels = v; |
| 18 | } |
| 19 | } |
| 20 | return features; |
| 21 | } |
| 22 | |
| 23 | VGGImpl::VGGImpl(std::vector<int> &cfg, int num_classes, bool batch_norm){ |
| 24 | features_ = make_features(cfg,batch_norm); |
no test coverage detected