Perceptual Loss & Feature Matching

1
#include <torch/torch.h>
2
#include <iostream>
3
#include <vector>
4
5
// A simple CNN feature extractor to mimic VGG-like layers
6
struct FeatureExtractorImpl : torch::nn::Module {
7
    // Convolutional layers
8
    torch::nn::Conv2d conv1{nullptr}, conv2{nullptr}, conv3{nullptr};
9
    FeatureExtractorImpl() {
10
        conv1 = register_module("conv1", torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 16, 3).stride(1).padding(1)));
11
        conv2 = register_module("conv2", torch::nn::Conv2d(torch::nn::Conv2dOptions(16, 32, 3).stride(1).padding(1)));
12
        conv3 = register_module("conv3", torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 3).stride(1).padding(1)));
13
    }
14
15
    // Forward returning intermediate feature maps after ReLU
16
    std::vector<torch::Tensor> forward_feats(const torch::Tensor& x_in) {
17
        std::vector<torch::Tensor> feats;
18
        auto x = torch::relu(conv1->forward(x_in));   // Layer 0 features (low-level edges)
19
        feats.push_back(x);
20
        x = torch::max_pool2d(x, 2);                  // Downsample
21
22
        x = torch::relu(conv2->forward(x));           // Layer 1 features (textures)
23
        feats.push_back(x);
24
        x = torch::max_pool2d(x, 2);
25
26
        x = torch::relu(conv3->forward(x));           // Layer 2 features (higher-level parts)
27
        feats.push_back(x);
28
        return feats;                                  // {B, C_l, H_l, W_l}
29
    }
30
};
31
TORCH_MODULE(FeatureExtractor);
32
33
// Compute weighted L2 perceptual loss across selected layers
34
// L = sum_l w_l * MSE(phi_l(x), phi_l(yhat))
35
torch::Tensor perceptual_loss_l2(const std::vector<torch::Tensor>& feats_x,
36
                                 const std::vector<torch::Tensor>& feats_yhat,
37
                                 const std::vector<double>& weights) {
38
    TORCH_CHECK(feats_x.size() == feats_yhat.size(), "Mismatched feature vector sizes");
39
    TORCH_CHECK(weights.size() == feats_x.size(), "Weights must match number of layers");
40
    torch::Tensor total = torch::zeros({}, feats_x[0].options());
41
    for (size_t l = 0; l < feats_x.size(); ++l) {
42
        auto diff = feats_x[l] - feats_yhat[l];                  // {B, C, H, W}
43
        auto mse = torch::mean(diff * diff);                     // scalar
44
        total = total + weights[l] * mse;                        // accumulate weighted layer loss
45
    }
46
    return total;                                                // scalar tensor
47
}
48
49
int main() {
50
    torch::manual_seed(42);
51
52
    // Hyperparameters and input shapes
53
    const int B = 2;      // batch size
54
    const int C = 3;      // RGB channels
55
    const int H = 128;    // height
56
    const int W = 128;    // width
57
58
    // Create random target image x and generated image yhat
59
    auto x = torch::rand({B, C, H, W});
60
    auto yhat = torch::rand({B, C, H, W});
61
62
    // Build feature extractor (in practice, often pretrained and frozen)
63
    FeatureExtractor phi;
64
    phi->eval(); // we're only extracting features here
65
66
    // Optional: normalize inputs if the extractor expects specific stats
67
    // For demo, we skip normalization; real-world use VGG/ImageNet mean-std.
68
69
    // Extract features for x and yhat
70
    auto feats_x = phi->forward_feats(x);
71
    auto feats_y = phi->forward_feats(yhat);
72
73
    // Example layer weights: emphasize shallow layers slightly for sharpness
74
    std::vector<double> weights = {1.0, 0.5, 0.25};
75
76
    // Compute perceptual loss
77
    auto L_perc = perceptual_loss_l2(feats_x, feats_y, weights);
78
79
    std::cout << "Perceptual L2 loss: " << L_perc.item<double>() << std::endl;
80
81
    // Backprop example: pretend yhat comes from a generator; here we just compute grad w.r.t yhat
82
    yhat.requires_grad_(true);
83
    auto feats_y_grad = phi->forward_feats(yhat);
84
    auto L_perc_grad = perceptual_loss_l2(feats_x, feats_y_grad, weights);
85
    L_perc_grad.backward();
86
    std::cout << "Grad computed through feature extractor: yes (" << yhat.grad().defined() << ")\n";
87
88
    return 0;
89
}
90

1
#include <torch/torch.h>
2
#include <iostream>
3
#include <vector>
4
5
// Simple discriminator that also exposes intermediate features
6
struct DiscriminatorImpl : torch::nn::Module {
7
    torch::nn::Conv2d conv1{nullptr}, conv2{nullptr}, conv3{nullptr};
8
    torch::nn::Linear fc{nullptr};
9
10
    DiscriminatorImpl() {
11
        conv1 = register_module("conv1", torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 32, 4).stride(2).padding(1)));
12
        conv2 = register_module("conv2", torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 4).stride(2).padding(1)));
13
        conv3 = register_module("conv3", torch::nn::Conv2d(torch::nn::Conv2dOptions(64, 128, 4).stride(2).padding(1)));
14
        fc    = register_module("fc",    torch::nn::Linear(128 * 16 * 16, 1)); // assumes input 128x128
15
    }
16
17
    // Forward that returns logits and a vector of features
18
    std::pair<torch::Tensor, std::vector<torch::Tensor>> forward_feats(const torch::Tensor& x_in) {
19
        std::vector<torch::Tensor> feats;
20
        auto x = torch::leaky_relu(conv1->forward(x_in), 0.2);
21
        feats.push_back(x);
22
        x = torch::leaky_relu(conv2->forward(x), 0.2);
23
        feats.push_back(x);
24
        x = torch::leaky_relu(conv3->forward(x), 0.2);
25
        feats.push_back(x);
26
        auto flat = x.view({x.size(0), -1});
27
        auto logit = fc->forward(flat);
28
        return {logit, feats};
29
    }
30
};
31
TORCH_MODULE(Discriminator);
32
33
// Feature matching: L_FM = sum_l || mean_{b,h,w}(f_l(real)) - mean_{b,h,w}(f_l(fake)) ||_1
34
torch::Tensor feature_matching_loss_L1(const std::vector<torch::Tensor>& feats_real,
35
                                       const std::vector<torch::Tensor>& feats_fake) {
36
    TORCH_CHECK(feats_real.size() == feats_fake.size(), "Mismatched feature vector sizes");
37
    torch::Tensor total = torch::zeros({}, feats_real[0].options());
38
    for (size_t l = 0; l < feats_real.size(); ++l) {
39
        // Reduce over batch and spatial dims: mean per-channel vector
40
        auto mr = feats_real[l].mean({0, 2, 3}); // {C}
41
        auto mf = feats_fake[l].mean({0, 2, 3}); // {C}
42
        auto l1 = torch::mean(torch::abs(mr - mf)); // scalar
43
        total = total + l1;
44
    }
45
    return total; // scalar tensor
46
}
47
48
int main() {
49
    torch::manual_seed(0);
50
51
    const int B = 4;  // batch size
52
    const int C = 3;  // channels
53
    const int H = 128, W = 128;
54
55
    // Real images (from dataset) and fake images (from generator); here we simulate both
56
    auto real = torch::rand({B, C, H, W});
57
    auto fake = torch::rand({B, C, H, W});
58
59
    Discriminator D;
60
    D->eval(); // when computing FM for G update, D's parameters are typically frozen
61
62
    // Forward both through D to collect features
63
    auto [logit_r, feats_r] = D->forward_feats(real.detach()); // detach to avoid grads into real path
64
    auto [logit_f, feats_f] = D->forward_feats(fake);          // grads will flow to 'fake' (generator)
65
66
    // Compute L1 feature matching loss across all exposed layers
67
    auto L_fm = feature_matching_loss_L1(feats_r, feats_f);
68
69
    std::cout << "Feature Matching L1 loss: " << L_fm.item<double>() << std::endl;
70
71
    // Backprop to show gradients can flow to 'fake' (as if from generator)
72
    fake.requires_grad_(true);
73
    auto [_, feats_f2] = D->forward_feats(fake);
74
    auto L_fm2 = feature_matching_loss_L1(feats_r, feats_f2);
75
    L_fm2.backward();
76
    std::cout << "Grad defined for fake: " << fake.grad().defined() << std::endl;
77
78
    return 0;
79
}
80