Initial commit

include/softmaxtree.cuh

@@ -0,0 +1,144 @@
+/* 
+ * File:   softmaxtree.h
+ * Author: Alex Krizhevsky
+ *
+ * Created on September 9, 2012, 5:50 PM
+ */
+
+#ifndef SOFTMAXTREE_H
+#define	SOFTMAXTREE_H
+
+#include <helper_cuda.h>
+#include <string>
+#include <map>
+#include <vector>
+#include <algorithm>
+#include <assert.h>
+
+#include <nvmatrix.cuh>
+#include <matrix.h>
+
+class SoftmaxNode;
+class SoftmaxTree;
+typedef std::vector<SoftmaxNode*> SoftmaxNodeV;
+
+class SoftmaxNode {
+    friend class SoftmaxTree;
+protected:
+    SoftmaxNodeV _children;
+    SoftmaxNode* _parent;
+    int _depth, _height, _size;
+    int _label;
+
+    /*
+     * Computes height for entire subtree rooted at this node and populates
+     * given height->nodes map.
+     */
+    int setDistances(std::map<int, SoftmaxNodeV*>& nodeHeights,
+                     std::map<int, SoftmaxNodeV*>& nodeDepths);
+    
+    void setNodeCounts(int &nodes, int& leaves);
+    /*
+     * Compute the number of leaves in this subtree, which is a good estimate
+     * of the number of training cases it represents.
+     */
+    int setSizes(ushort* nodeSizes);
+    
+public:
+    SoftmaxNode(SoftmaxNode* parent, int label);
+    ~SoftmaxNode();
+    SoftmaxNode& addChild(int label);
+    
+    int getDepth() const;
+    int getHeight() const;
+    int getLabel() const;
+    int getSize() const;
+    SoftmaxNode* getParent(); // Might be null, so must be pointer
+    SoftmaxNodeV& getChildren();
+};
+
+
+/*
+ * numLabels: the number of leaves in the tree (normally 1000)
+ * numNodes: the total number of nodes in the tree
+ */
+class SoftmaxTree {
+    friend class SoftmaxNode;
+protected:
+    SoftmaxNode* _root;
+    std::map<int, SoftmaxNodeV*> _nodeHeights, _nodeDepths;
+    /*
+     * Map from depth --> ushort2[]
+     * where each ushort2 gives the index and parent index
+     * of a node at the given depth.
+     */
+    std::map<int, ushort2*> _nodeFwdMeta;
+    /*
+     * Map from height --> ushort2[]
+     * where each ushort2 gives the index and number of children
+     * of a node at the given height.
+     */
+    std::map<int, ushort2*> _nodeBwdMeta;
+    /*
+     * Map from height --> ushort[][]
+     * where each ushort[] gives children of a given node at a given height.
+     */
+    std::map<int, ushort**> _nodeChildMeta;
+    
+    /*
+     * An array of length numNodes with index i storing the number
+     * of leaves in subtree rooted at node with label i.
+     */
+    ushort* _nodeSizes;
+    int _numNodes, _numLeaves;
+    void setDistances();
+    void setNodeCounts();
+    void setNodeSizes();
+    void setFwdMeta();
+    void setBwdMeta();
+    void preprocess(NVMatrix& inp);
+    void postprocess(NVMatrix& inp);
+public:
+    SoftmaxTree(int rootLabel);
+    ~SoftmaxTree();
+    
+    void finalize();
+    
+    SoftmaxNode& getRoot();
+    SoftmaxNodeV& getNodesAtHeight(int height);
+    SoftmaxNodeV& getNodesAtDepth(int depth);
+    int getHeight() const;
+    int getDepth() const;
+    int getNumLeaves() const;
+    int getNumNodes() const;
+    
+    /*
+     * offsets: (numNodes, numFeatures)
+     * targets: (numNodes, numFeatures) 
+     */
+    void makeWeights(NVMatrix& offsets, NVMatrix& targets);
+    
+    /*
+     * grads: (numNodes, numFeatures)
+     * 
+     * The idea is that grads contains gradients for the leaves 
+     * (i.e. the first numLabels rows), so this routine will
+     * distribute them up the tree.
+     */
+    void distributeGradients(NVMatrix& grads);
+    
+    /*
+     * inc := mom * inc - wc * epsW * weight + epsW * grad
+     * weight := weight + inc
+     * 
+     * weights: (numNodes, numFeatures)
+     * incs:    (numNodes, numFeatures)
+     * grads:   (numNodes , numFeatures)
+     */
+    void updateWeights(NVMatrix& weights, NVMatrix& incs, NVMatrix& grads, float epsWBase, float mom, float wcBase);
+    
+    
+};
+
+#endif	/* SOFTMAXTREE_H */
+