void CosSimVecMatLayer::backward(const UpdateCallback& callback) {
CHECK_EQ(backward_.size(), 1UL) << "Only one forward function needed";
MatrixPtr inV0 = getInputValue(0);
MatrixPtr inV1 = getInputValue(1);
MatrixPtr inG0 = getInputGrad(0);
MatrixPtr inG1 = getInputGrad(1);
MatrixPtr outV = getOutputValue();
MatrixPtr outG = getOutputGrad();
size_t batchSize = inV0->getHeight();
CHECK(inV0 && inV1 && inG0 && inG1 && outV && outG);
REGISTER_TIMER_INFO("BwCosVMTimer", getName().c_str());
for (size_t i = 0; i < batchSize; i++) {
tmpRow0->setData(inV0->rowBuf(i));
tmpRow1->setData(inG0->rowBuf(i));
tmpMtx0->setData(inV1->rowBuf(i));
tmpMtx1->setData(inG1->rowBuf(i));
tmpRow2->setData(outV->rowBuf(i));
tmpRow3->setData(outG->rowBuf(i));
BufferArgs inputs;
BufferArgs outputs;
inputs.addArg(*tmpRow3);
inputs.addArg(*tmpRow2);
inputs.addArg(*tmpMtx0);
inputs.addArg(*tmpRow0);
outputs.addArg(*tmpMtx1, ADD_TO);
outputs.addArg(*tmpRow1, ADD_TO);
backward_[0]->calc(inputs, outputs);
}
}
void CosSimVecMatLayer::forward(PassType passType) {
Layer::forward(passType);
CHECK_EQ(forward_.size(), 1UL) << "Only one forward function needed";
MatrixPtr inV0 = getInputValue(0);
MatrixPtr inV1 = getInputValue(1);
size_t batchSize = inV0->getHeight();
size_t numKeys = getSize();
CHECK_EQ(batchSize, inV1->getHeight());
{
REGISTER_TIMER_INFO("FwResetTimer", getName().c_str());
reserveOutput(batchSize, numKeys);
}
MatrixPtr outV = getOutputValue();
CHECK(outV && inV0 && inV1);
REGISTER_TIMER_INFO("FwCosVMTimer", getName().c_str());
for (size_t i = 0; i < batchSize; i++) {
tmpRow0->setData(inV0->rowBuf(i));
tmpMtx0->setData(inV1->rowBuf(i));
tmpRow2->setData(outV->rowBuf(i));
BufferArgs inputs;
BufferArgs outputs;
inputs.addArg(*tmpMtx0);
inputs.addArg(*tmpRow0);
outputs.addArg(*tmpRow2, ASSIGN_TO);
forward_[0]->calc(inputs, outputs);
}
}
