# 目录

### 一、添加Loss Function

#### 1. gan_loss_layer.hpp

``````//  Create on: 2016/10/19 ShanghaiTech
//  Author:    Yingying Zhang

#ifndef CAFFE_GAN_LOSS_LAYER_HPP_
#define CAFFE_GAN_LOSS_LAYER_HPP_

#include <vector>

#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"

#include "caffe/layers/loss_layer.hpp"

namespace caffe {

/*This loss layer computes loss for GAN.
*  1. loss for discriminator
*     bottom[0] : D(x)
*     bottom[1] : D(G(z))
*  2. loss for generator
*     bottom[0] : D(G(z))
*/

template <typename Dtype>
class GANLossLayer : public LossLayer<Dtype> {
public:
explicit GANLossLayer(const LayerParameter& param)
: LossLayer<Dtype>(param)  {}
virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Reshape(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
vector<int> loss_shape(0);  // Loss layers output a scalar; 0 axes.
top[0]->Reshape(loss_shape);
}
virtual inline int ExactNumBottomBlobs() const { return -1; }
virtual inline int MinBottomBlobs() const { return 1; }
virtual inline int MaxBottomBlobs() const { return 2; }
virtual inline const char* type() const { return "GANLoss"; }
virtual inline bool AllowForceBackward(const int bottom_index) const {
return true;
}

protected:
virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

int dis_iter_;
int gen_iter_;
int iter_idx_;
};

template <typename Dtype>
class GANDGLossLayer : public LossLayer<Dtype> {
public:
explicit GANDGLossLayer(const LayerParameter& param)
: LossLayer<Dtype>(param)  {}
virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Reshape(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
vector<int> loss_shape(0);  // Loss layers output a scalar; 0 axes.
top[0]->Reshape(loss_shape);
}
virtual inline int ExactNumBottomBlobs() const { return 1; }
virtual inline const char* type() const { return "GANDGLoss"; }
virtual inline bool AllowForceBackward(const int bottom_index) const {
return true;
}

protected:
virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);

int dis_iter_;
int gen_iter_;
int gan_mode_;
int diter_idx_;
int giter_idx_;
};

}  // namespace caffe

#endif  // CAFFE_GAN_LOSS_LAYER_HPP_
``````

#### 1. gan_loss_layer.cpp

``````//  Create on: 2016/10/19 ShanghaiTech
//  Author:    Yingying Zhang

#include <algorithm>
#include <vector>

#include "caffe/layers/gan_loss_layer.hpp"

namespace caffe {

template <typename Dtype>
void GANLossLayer<Dtype>::LayerSetUp(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
LossLayer<Dtype>::LayerSetUp(bottom, top);
iter_idx_ = 0;
dis_iter_ = this->layer_param_.gan_loss_param().dis_iter();
gen_iter_ = this->layer_param_.gan_loss_param().gen_iter();
//discriminative mode
if (bottom.size() == 2) {
CHECK_EQ(bottom[0]->shape(0), bottom[1]->shape(0));
CHECK_EQ(bottom[0]->shape(1), 1);
CHECK_EQ(bottom[1]->shape(1), 1);
}
//generative mode
if (bottom.size() == 1) {
CHECK_EQ(bottom[0]->shape(1), 1);
}
}

template <typename Dtype>
void GANLossLayer<Dtype>::Forward_cpu(
const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
int batch_size = bottom[0]->count();
Dtype loss(0.0);
//1. discriminative mode
if (bottom.size() == 2) {
const Dtype* score1 = bottom[0]->cpu_data();
const Dtype* score2 = bottom[1]->cpu_data();
for(int i = 0; i<batch_size; ++i) {
loss -= std::log(score1[i]) + std::log(1 - score2[i]);
}
}
//2. generative mode
if (bottom.size() == 1) {
const Dtype* score = bottom[0]->cpu_data();
for(int i = 0; i<batch_size; ++i) {
loss -= std::log(score[i]);
}
}
loss /= static_cast<Dtype>(batch_size);
top[0]->mutable_cpu_data()[0] = loss;
iter_idx_++;
}

template <typename Dtype>
void GANLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
int batch_size = bottom[0]->count();
//1. discriminative mode
if (bottom.size() == 2) {
if (iter_idx_ % dis_iter_ == 0) {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(-1) /
bottom[0]->cpu_data()[i] / static_cast<Dtype>(batch_size);
bottom[1]->mutable_cpu_diff()[i] = Dtype(-1) /
(bottom[1]->cpu_data()[i] - Dtype(1))  / static_cast<Dtype>(batch_size);
}
} else {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(0);
bottom[1]->mutable_cpu_diff()[i] = Dtype(0);
}
}
}
//2. generative mode
if (bottom.size() == 1) {
if (iter_idx_ % gen_iter_ == 0) {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(-1) /
bottom[0]->cpu_data()[i] / static_cast<Dtype>(batch_size);
}
} else {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(0);
}
}
}
}

INSTANTIATE_CLASS(GANLossLayer);
REGISTER_LAYER_CLASS(GANLoss);

template <typename Dtype>
void GANDGLossLayer<Dtype>::LayerSetUp(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
LossLayer<Dtype>::LayerSetUp(bottom, top);
diter_idx_ = 0;
giter_idx_ = 0;
dis_iter_ = this->layer_param_.gan_loss_param().dis_iter();
gen_iter_ = this->layer_param_.gan_loss_param().gen_iter();
gan_mode_ = 1;
}

template <typename Dtype>
void GANDGLossLayer<Dtype>::Forward_cpu(
const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
int batch_size = bottom[0]->count();
const Dtype* score = bottom[0]->cpu_data();
Dtype loss(0.0);
//when gan_mode_ = 1, the input of loss is D(x)
//loss is discriminative loss: -log(D(x))
if (gan_mode_ == 1) {
diter_idx_++;
for(int i = 0; i<batch_size; ++i) {
loss -= std::log(score[i]);
}
}
//when gan_mode_ = 2, the input of loss is D(G(z))
//loss is discriminative loss: -log(1-D(G(z)))
if (gan_mode_ == 2){
for(int i = 0; i<batch_size; ++i) {
loss -= std::log(1 - score[i]);
}
}
//when gan_mode_ = 3, the input of loss is D(G(z))
//loss is generative loss: -log(D(G(z)))
if (gan_mode_ == 3){
giter_idx_++;
for(int i = 0; i<batch_size; ++i) {
loss -= std::log(score[i]);
}
}
loss /= static_cast<Dtype>(batch_size);
top[0]->mutable_cpu_data()[0] = loss;
}

template <typename Dtype>
void GANDGLossLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
int batch_size = bottom[0]->count();
//when gan_mode_ = 1, the input of loss is D(x)
//backward for discriminative loss
if (gan_mode_ == 1) {
if (diter_idx_ % dis_iter_ == 0 ) {
for(int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(-1) /
bottom[0]->cpu_data()[i] / static_cast<Dtype>(batch_size);
}
} else {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(0);
}
}
}
//when gan_mode_ = 2, the input of loss is D(G(z))
//backward for discriminative loss
if (gan_mode_ == 2){
if (diter_idx_ % dis_iter_ == 0 ) {
for(int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(-1) /
(bottom[0]->cpu_data()[i] - Dtype(1))  / static_cast<Dtype>(batch_size);
}
} else {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(0);
}
}
}
//when gan_mode_ = 3, the input of loss is D(G(z))
//backward for generative loss
if (gan_mode_ == 3){
if (giter_idx_ % gen_iter_ == 0 ) {
for(int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(-1) /
bottom[0]->cpu_data()[i] / static_cast<Dtype>(batch_size);
}
} else {
for (int i = 0; i<batch_size; ++i) {
bottom[0]->mutable_cpu_diff()[i] = Dtype(0);
}
}
}
// update gan_mode_
gan_mode_ = gan_mode_ == 3 ? 1 : gan_mode_ + 1;

}

INSTANTIATE_CLASS(GANDGLossLayer);
REGISTER_LAYER_CLASS(GANDGLoss);

}  // namespace caffe
``````

### 二、添加gan_gate_layer

#### gan_gate_layer.hpp

``````//  Create on: 2016/10/19 ShanghaiTech
//  Author:    Yingying Zhang

#ifndef CAFFE_GAN_GATE_LAYER_HPP_
#define CAFFE_GAN_GATE_LAYER_HPP_

#include <vector>
#include "caffe/blob.hpp"
#include "caffe/layer.hpp"

namespace caffe {

template <typename Dtype>
class GANGateLayer : public Layer<Dtype> {
public:
explicit GANGateLayer(const LayerParameter& param)
: Layer<Dtype>(param) {}
virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {}

virtual inline const char* type() const { return "GANGate"; }
virtual inline int ExactNumBottomBlobs() const { return 2; }
virtual inline int ExactNumTopBlobs() const { return 1; }

protected:
virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {}
int gan_mode_;
};

}  // namespace caffe

#endif  // CAFFE_GAN_GATE_LAYER_HPP_
``````

#### gan_gate_layer.cpp

``````//  Create on: 2016/10/21 ShanghaiTech
//  Author:    Yingying Zhang

#include <vector>
#include "caffe/layers/gan_gate_layer.hpp"

namespace caffe {

template <typename Dtype>
void GANGateLayer<Dtype>::LayerSetUp(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
gan_mode_ = 1;
top[0]->ReshapeLike(*bottom[0]);
}

template <typename Dtype>
void GANGateLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
CHECK_EQ(bottom[0]->num(), bottom[1]->num());
CHECK_EQ(bottom[0]->channels(), bottom[1]->channels());
CHECK_EQ(bottom[0]->height(), bottom[1]->height());
CHECK_EQ(bottom[0]->width(), bottom[1]->width());
int index = gan_mode_ == 1 ? 0 : 1;
top[0]->ReshapeLike(*bottom[index]);
top[0]->ShareData(*bottom[index]);
top[0]->ShareDiff(*bottom[index]);
gan_mode_ = gan_mode_ == 3 ? 1 : gan_mode_ + 1;
}

INSTANTIATE_CLASS(GANGateLayer);
REGISTER_LAYER_CLASS(GANGate);

}  // namespace caffe``````

### 三、添加rand_vec_layer

#### rand_vec_layer.hpp

``````//  Create on: 2016/10/19 ShanghaiTech
//  Author:    Yingying Zhang

#ifndef CAFFE_RAND_VEC_LAYER_HPP_
#define CAFFE_RAND_VEC_LAYER_HPP_

#include <vector>

#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"
#include "caffe/util/rng.hpp"
namespace caffe {

template <typename Dtype>
class RandVecLayer : public Layer<Dtype> {
public:
explicit RandVecLayer(const LayerParameter& param)
: Layer<Dtype>(param) {}
virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {}

virtual inline const char* type() const { return "RandVec"; }
virtual inline int ExactNumBottomBlobs() const { return 0; }
virtual inline int ExactNumTopBlobs() const { return 1; }

protected:
virtual Dtype GetRandom(const Dtype lower, const Dtype upper);
virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top);
virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {}

shared_ptr<Caffe::RNG> data_rng_;
int batch_size_;
int dim_;
Dtype lower_;
Dtype upper_;
int iter_idx_;
int repeat_;
int height_;
int width_;
};

}  // namespace caffe

#endif  // CAFFE_RAND_VEC_LAYER_HPP_
``````

#### rand_vec_layer.cpp

``````//  Create on: 2016/10/19 ShanghaiTech
//  Author:    Yingying Zhang

#include <vector>

#include "caffe/layer.hpp"
#include "caffe/layers/rand_vec_layer.hpp"

#define MAX_RANDOM 10000

namespace caffe {

template <typename Dtype>
Dtype RandVecLayer<Dtype>::GetRandom(const Dtype lower, const Dtype upper) {
CHECK(data_rng_);
CHECK_LT(lower, upper) << "Upper bound must be greater than lower bound!";
caffe::rng_t* data_rng =
static_cast<caffe::rng_t*>(data_rng_->generator());
return static_cast<Dtype>(((*data_rng)()) % static_cast<unsigned int>(
(upper - lower) * MAX_RANDOM)) / static_cast<Dtype>(MAX_RANDOM)+lower;
}

template <typename Dtype>
void RandVecLayer<Dtype>::LayerSetUp(
const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
const RandVecParameter& rand_vec_param = this->layer_param_.rand_vec_param();
batch_size_ = rand_vec_param.batch_size();
dim_ = rand_vec_param.dim();
height_ = rand_vec_param.height();
width_ = rand_vec_param.width();
lower_ = rand_vec_param.lower();
upper_ = rand_vec_param.upper();
repeat_ = rand_vec_param.repeat();
iter_idx_ = 1;
vector<int> top_shape(2);
top_shape[0] = batch_size_;
top_shape[1] = dim_;
if (height_ >0 && width_>0) {
top_shape.resize(4);
top_shape[0] = batch_size_;
top_shape[1] = dim_;
top_shape[2] = height_;
top_shape[3] = width_;
}
top[0]->Reshape(top_shape);
}

template <typename Dtype>
void RandVecLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
if (iter_idx_ == 1) {
const unsigned int data_rng_seed = caffe_rng_rand();
data_rng_.reset(new Caffe::RNG(data_rng_seed));
int count = top[0]->count();
for (int i = 0; i<count; ++i) {
top[0]->mutable_cpu_data()[i] = GetRandom(lower_, upper_);
}
}
//update iter_idx_
iter_idx_ = iter_idx_ == repeat_ ? 1 : iter_idx_ + 1;
}

INSTANTIATE_CLASS(RandVecLayer);
REGISTER_LAYER_CLASS(RandVec);

}  // namespace caffe
``````

### 四、修改Solver中Step函数

``````template <typename Dtype>
void Solver<Dtype>::Step(int iters) {
const int start_iter = iter_;
const int stop_iter = iter_ + iters;
int average_loss = this->param_.average_loss();
bool gan_solver = this->param_.gan_solver();
losses_.clear();
smoothed_loss_ = 0;
Dtype loss_buff = 0;
int gan_iter_times = 2;

while (iter_ < stop_iter) {
// zero-init the params
net_->ClearParamDiffs();
if (param_.test_interval() && iter_ % param_.test_interval() == 0
&& (iter_ > 0 || param_.test_initialization())
&& Caffe::root_solver()) {
TestAll();
if (requested_early_exit_) {
// Break out of the while loop because stop was requested while testing.
break;
}
}

for (int i = 0; i < callbacks_.size(); ++i) {
callbacks_[i]->on_start();
}
const bool display = param_.display() && iter_ % param_.display() == 0;
net_->set_debug_info(display && param_.debug_info());
// accumulate the loss and gradient
Dtype loss = 0;
if (gan_solver) {
for (int i = 0; i < gan_iter_times; ++i) {
loss += net_->ForwardBackward();
}
gan_iter_times = gan_iter_times == 2 ? 1 : 2;
} else {
for (int i = 0; i < param_.iter_size(); ++i) {
loss += net_->ForwardBackward();
}
}
loss /= param_.iter_size();

// average the loss across iterations for smoothed reporting
if (display) {
LOG_IF(INFO, Caffe::root_solver()) << "Iteration " << iter_
<< ", loss = " << smoothed_loss_;
if (gan_solver) {
if (gan_iter_times == 1 && iter_ != 0) {
LOG_IF(INFO, Caffe::root_solver()) << "    dis_loss =  "  << loss;
LOG_IF(INFO, Caffe::root_solver()) << "    gen_loss =  "  << loss_buff;
}
}
const vector<Blob<Dtype>*>& result = net_->output_blobs();
int score_index = 0;
for (int j = 0; j < result.size(); ++j) {
const Dtype* result_vec = result[j]->cpu_data();
const string& output_name =
net_->blob_names()[net_->output_blob_indices()[j]];
const Dtype loss_weight =
net_->blob_loss_weights()[net_->output_blob_indices()[j]];
for (int k = 0; k < result[j]->count(); ++k) {
ostringstream loss_msg_stream;
if (loss_weight) {
loss_msg_stream << " (* " << loss_weight
<< " = " << loss_weight * result_vec[k] << " loss)";
}
LOG_IF(INFO, Caffe::root_solver()) << "    Train net output #"
<< score_index++ << ": " << output_name << " = "
<< result_vec[k] << loss_msg_stream.str();
}
}
}
loss_buff = loss;
for (int i = 0; i < callbacks_.size(); ++i) {
}
ApplyUpdate();

// Increment the internal iter_ counter -- its value should always indicate
// the number of times the weights have been updated.
++iter_;

SolverAction::Enum request = GetRequestedAction();

// Save a snapshot if needed.
if ((param_.snapshot()
&& iter_ % param_.snapshot() == 0
&& Caffe::root_solver()) ||
(request == SolverAction::SNAPSHOT)) {
Snapshot();
}
if (SolverAction::STOP == request) {
requested_early_exit_ = true;
// Break out of training loop.
break;
}
}
}``````

PS：看到很多同学索要全部修改过的源码，现在给出下载链接：GAN相关源码