generator_population.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <memory>
#include <stdexcept>
#include <vector>
 
#include "pulsatrix/module.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class GeneratorPopulation {
public:
    explicit GeneratorPopulation(std::vector<std::unique_ptr<Module>> generators)
        : generators_(std::move(generators)) {
        if (generators_.empty()) {
            throw std::invalid_argument("GeneratorPopulation: must have at least one generator");
        }
        for (const auto& g : generators_) {
            if (g == nullptr) {
                throw std::invalid_argument("GeneratorPopulation: generator entries must not be null");
            }
        }
    }
 
    [[nodiscard]] size_t size() const { return generators_.size(); }
 
    [[nodiscard]] Module& generator(size_t index) { return *generators_.at(index); }
    [[nodiscard]] const Module& generator(size_t index) const { return *generators_.at(index); }
 
    void ReplaceGenerator(size_t index, std::unique_ptr<Module> new_generator) {
        if (new_generator == nullptr) {
            throw std::invalid_argument("GeneratorPopulation::ReplaceGenerator: new_generator must not be null");
        }
        generators_.at(index) = std::move(new_generator);
    }
 
private:
    std::vector<std::unique_ptr<Module>> generators_;
};
 
inline Tensor SliceBatch(const Tensor& t, int64_t start, int64_t count, DeviceBackend* backend) {
    if (t.shape().rank() < 1) {
        throw std::invalid_argument("SliceBatch: t must have rank >= 1");
    }
    int64_t leading = t.shape().dim(0);
    if (count <= 0 || start < 0 || start + count > leading) {
        throw std::invalid_argument("SliceBatch: [start, start+count) must be a valid sub-range of t's leading dimension");
    }
 
    std::vector<int64_t> out_dims{count};
    for (int64_t i = 1; i < t.shape().rank(); ++i) {
        out_dims.push_back(t.shape().dim(i));
    }
    Shape out_shape(out_dims);
 
    int64_t row_size = (leading > 0) ? (t.numel() / leading) : 0;
    std::vector<float> values(static_cast<size_t>(row_size * count));
    for (int64_t i = 0; i < row_size * count; ++i) {
        values[static_cast<size_t>(i)] = t.data()[start * row_size + i];
    }
    return Tensor(out_shape, backend, values, t.device());
}
 
inline std::vector<float> FlattenParameters(Module& module) {
    std::vector<float> flat;
    for (const auto& p : module.parameters()) {
        for (int64_t i = 0; i < p.value->numel(); ++i) {
            flat.push_back(p.value->data()[i]);
        }
    }
    return flat;
}
 
inline void RestoreParameters(Module& module, const std::vector<float>& flat) {
    size_t offset = 0;
    for (const auto& p : module.parameters()) {
        int64_t n = p.value->numel();
        for (int64_t i = 0; i < n; ++i) {
            if (offset >= flat.size()) {
                throw std::invalid_argument("RestoreParameters: flat has fewer entries than module.parameters() needs");
            }
            p.value->data()[i] = flat[offset];
            ++offset;
        }
    }
    if (offset != flat.size()) {
        throw std::invalid_argument("RestoreParameters: flat has more entries than module.parameters() needs");
    }
}
 
inline void ZeroModuleGradients(Module& module) {
    for (const auto& p : module.parameters()) {
        for (int64_t i = 0; i < p.grad->numel(); ++i) {
            p.grad->data()[i] = 0.0f;
        }
    }
}
 
inline Tensor GeneratePooledFakeSamples(GeneratorPopulation& population,
                                         const std::vector<Tensor>& noise_per_generator, DeviceBackend* backend) {
    if (noise_per_generator.size() != population.size()) {
        throw std::invalid_argument("GeneratePooledFakeSamples: noise_per_generator.size() must equal population.size()");
    }
    std::vector<Tensor> fakes;
    fakes.reserve(population.size());
    for (size_t i = 0; i < population.size(); ++i) {
        fakes.push_back(population.generator(i).forward(noise_per_generator[i]));
    }
    return Tensor::Stack(fakes, backend);
}
 
inline void BackwardThroughPopulation(GeneratorPopulation& population, const Tensor& pooled_grad,
                                       const std::vector<int64_t>& batch_sizes, DeviceBackend* backend) {
    if (batch_sizes.size() != population.size()) {
        throw std::invalid_argument("BackwardThroughPopulation: batch_sizes.size() must equal population.size()");
    }
    int64_t total = 0;
    for (int64_t b : batch_sizes) {
        total += b;
    }
    if (pooled_grad.shape().rank() < 1 || total != pooled_grad.shape().dim(0)) {
        throw std::invalid_argument("BackwardThroughPopulation: batch_sizes must sum to pooled_grad's leading dimension");
    }
 
    int64_t offset = 0;
    for (size_t i = 0; i < population.size(); ++i) {
        int64_t count = batch_sizes[i];
        Tensor grad_slice = SliceBatch(pooled_grad, offset, count, backend);
        (void)population.generator(i).backward(grad_slice);
        offset += count;
    }
}
 
}  // namespace pulsatrix