Function pretrained_resnet

lesson6-Segmentation/ResNet.cpp:195–240

}

ResNet pretrained_resnet(int64_t num_classes, std::string model_name, std::string weight_path){
    std::map<std::string, std::vector<int>> name2layers = getParams();
	int groups = 1;
	int width_per_group = 64;
	if (model_name == "resnext50_32x4d") {
		groups = 32; width_per_group = 4;
	}
	if (model_name == "resnext101_32x8d") {
		groups = 32; width_per_group = 8;
	}
    ResNet net_pretrained = ResNet(name2layers[model_name],1000,model_name,groups,width_per_group);
    torch::load(net_pretrained, weight_path);
    if(num_classes == 1000) return net_pretrained;
    ResNet module = ResNet(name2layers[model_name],num_classes,model_name);

    torch::OrderedDict<std::string, at::Tensor> pretrained_dict = net_pretrained->named_parameters();
    torch::OrderedDict<std::string, at::Tensor> model_dict = module->named_parameters();

    for (auto n = pretrained_dict.begin(); n != pretrained_dict.end(); n++)
    {
        if (strstr((*n).key().data(), "fc.")) {
            continue;
        }
        model_dict[(*n).key()] = (*n).value();
    }

    torch::autograd::GradMode::set_enabled(false);  // make parameters copying possible
    auto new_params = model_dict; // implement this
    auto params = module->named_parameters(true /*recurse*/);
    auto buffers = module->named_buffers(true /*recurse*/);
    for (auto& val : new_params) {
        auto name = val.key();
        auto* t = params.find(name);
        if (t != nullptr) {
            t->copy_(val.value());
        }
        else {
            t = buffers.find(name);
            if (t != nullptr) {
                t->copy_(val.value());
            }
        }
    }
    torch::autograd::GradMode::set_enabled(true);
    return module;
}