ailayer_conv2d.h

Go to the documentation of this file.

 
#ifndef CONV2D_LAYER
#define CONV2D_LAYER
 
#include "core/aifes_core.h"
 
#define AIFES_CHANNELS_FIRST    1
#define AIFES_CHANNELS_LAST     -1
 
#define AIFES_PADDING_SAME      0xFFFF  // Not fully implemented yet
#define AIFES_PADDING_VALID     0
 
typedef struct ailayer_conv2d   ailayer_conv2d_t;
 
struct ailayer_conv2d {
    ailayer_t base; 
    // Values (configurable)
    uint32_t filter_count;  
    uint16_t kernel_size[2];    
    uint16_t stride[2];  
    uint16_t dilation[2];  
    uint16_t padding[2];  // AIFES_PADDING_SAME = 0xFFFF
 
    int8_t channel_axis; 
 
    aitensor_t weights; 
    aitensor_t bias;   
    uint16_t weights_shape[4]; 
    uint16_t bias_shape[1]; 
    aitensor_t *trainable_params[2]; 
    aitensor_t *gradients[2]; 
    void *optimem[2];   
 
    uint16_t result_shape[4]; 
 
    void (*conv2d_fwd)(
                    const aitensor_t *input,
                    const uint16_t stride[2],    // [s_h, s_w]
                    const uint16_t dilation[2],  // [d_h, d_w]
                    const uint16_t padding[2],   // AIFES_PADDING_SAME not allowed
                    const aitensor_t *weights,
                    const aitensor_t *bias,
                    int8_t channel_axis,
                    void *work_space,
                    aitensor_t *output
                    );
 
    void (*conv2d_bwd)(
                    const aitensor_t *input,
                    const uint16_t stride[2],    // [s_h, s_w]
                    const uint16_t dilation[2],  // [d_h, d_w]
                    const uint16_t padding[2],
                    const aitensor_t *output,
                    int8_t channel_axis,
                    void *work_space,
                    aitensor_t *d_weights
    );
 
    void (*conv2d_bwd_full)(
                    const aitensor_t *delta_out,
                    const uint16_t stride[2],    // [s_h, s_w]
                    const uint16_t dilation[2],  // [d_h, d_w]
                    const uint16_t padding[2],
                    const aitensor_t *weights,
                    int8_t channel_axis,
                    void *work_space,
                    aitensor_t *delta_in
    );
 
    void (*tensor_add)(const aitensor_t *a, const aitensor_t *b, aitensor_t *result);
 
    void (*sum_channelwise)(const aitensor_t *x, int8_t channel_axis, aitensor_t *result);
};
 
extern const aicore_layertype_t *ailayer_conv2d_type;
 
ailayer_t *ailayer_conv2d(ailayer_conv2d_t *layer, ailayer_t *input_layer);
 
void ailayer_conv2d_forward(ailayer_t *self);
 
void ailayer_conv2d_backward(ailayer_t *self);
 
void ailayer_conv2d_calc_result_shape(ailayer_t *self);
 
uint32_t ailayer_conv2d_sizeof_bwdmem(const ailayer_t *self);
 
uint32_t ailayer_conv2d_sizeof_paramem(const ailayer_t *self);
 
void ailayer_conv2d_set_paramem(ailayer_t *self, void *memory_ptr);
 
uint32_t ailayer_conv2d_sizeof_trainmem(const ailayer_t *self);
 
void ailayer_conv2d_set_trainmem(ailayer_t *self, void *memory_ptr);
 
#ifdef AIDEBUG_PRINT_MODULE_SPECS
void ailayer_conv2d_print_specs(const ailayer_t *self);
#endif // AIDEBUG_PRINT_MODULE_SPECS
 
#endif // CONV2D_LAYER