Initial implementation of residual neural networks

main
Shad Amethyst 2 years ago
parent fa0bc0be9f
commit 83dc763746

@ -122,6 +122,14 @@ impl<F, Act, Reg, R: Rng> NeuraDenseLayerPartial<F, Act, Reg, R> {
}
}
impl<F: Float, Act: NeuraDerivable<F>, Reg: NeuraDerivable<F>> NeuraShapedLayer
for NeuraDenseLayer<F, Act, Reg>
{
fn output_shape(&self) -> NeuraShape {
NeuraShape::Vector(self.weights.shape().0)
}
}
impl<
F: Float + std::fmt::Debug + 'static,
Act: NeuraDerivable<F>,
@ -144,10 +152,6 @@ where
self.regularization,
))
}
fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
NeuraShape::Vector(constructed.weights.shape().0)
}
}
impl<

@ -35,6 +35,12 @@ impl<R: Rng> NeuraDropoutLayer<R> {
}
}
impl<R: Rng> NeuraShapedLayer for NeuraDropoutLayer<R> {
fn output_shape(&self) -> NeuraShape {
self.shape
}
}
impl<R: Rng> NeuraPartialLayer for NeuraDropoutLayer<R> {
type Constructed = NeuraDropoutLayer<R>;
@ -45,10 +51,6 @@ impl<R: Rng> NeuraPartialLayer for NeuraDropoutLayer<R> {
self.mask = DVector::from_element(input_shape.size(), false);
Ok(self)
}
fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
constructed.shape
}
}
impl<R: Rng, F: Float> NeuraLayer<DVector<F>> for NeuraDropoutLayer<R> {

@ -48,13 +48,15 @@ impl<Input: Clone> NeuraLayer<Input> for () {
}
}
pub trait NeuraShapedLayer {
fn output_shape(&self) -> NeuraShape;
}
pub trait NeuraPartialLayer {
type Constructed;
type Constructed: NeuraShapedLayer;
type Err;
fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err>;
fn output_shape(constructed: &Self::Constructed) -> NeuraShape;
}
pub trait NeuraTrainableLayerBase<Input>: NeuraLayer<Input> {

@ -23,6 +23,12 @@ impl NeuraNormalizeLayer {
}
}
impl NeuraShapedLayer for NeuraNormalizeLayer {
fn output_shape(&self) -> NeuraShape {
self.shape
}
}
impl NeuraPartialLayer for NeuraNormalizeLayer {
type Constructed = NeuraNormalizeLayer;
@ -31,10 +37,6 @@ impl NeuraPartialLayer for NeuraNormalizeLayer {
fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
Ok(Self { shape: input_shape })
}
fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
constructed.shape
}
}
impl<F: Float + Scalar> NeuraLayer<DVector<F>> for NeuraNormalizeLayer {

@ -41,6 +41,12 @@ impl<F: Float + Scalar + NumAssignOps> NeuraLayer<DVector<F>> for NeuraSoftmaxLa
}
}
impl NeuraShapedLayer for NeuraSoftmaxLayer {
fn output_shape(&self) -> NeuraShape {
self.shape
}
}
impl NeuraPartialLayer for NeuraSoftmaxLayer {
type Constructed = Self;
type Err = ();
@ -48,10 +54,6 @@ impl NeuraPartialLayer for NeuraSoftmaxLayer {
fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
Ok(Self { shape: input_shape })
}
fn output_shape(constructed: &Self::Constructed) -> NeuraShape {
constructed.shape
}
}
impl<F: Float + Scalar + NumAssignOps> NeuraTrainableLayerBase<DVector<F>> for NeuraSoftmaxLayer {

@ -1,6 +1,7 @@
#![feature(generic_arg_infer)]
// #![feature(generic_const_exprs)]
#![feature(associated_type_defaults)]
#![feature(arc_unwrap_or_clone)]
pub mod algebra;
pub mod derivable;

@ -2,6 +2,7 @@ use crate::{
algebra::NeuraVectorSpace, gradient_solver::NeuraGradientSolverBase, layer::NeuraLayer,
};
pub mod residual;
pub mod sequential;
// TODO: extract regularize from this, so that we can drop the trait constraints on NeuraSequential's impl

@ -0,0 +1,63 @@
use std::borrow::Borrow;
use nalgebra::{Const, DVector, Dyn, VecStorage};
use crate::prelude::NeuraShape;
#[derive(Clone, Copy, Debug)]
pub struct NeuraAxisAppend;
#[derive(Clone, Copy, Debug)]
pub enum NeuraAxisErr {
NoInput,
ConflictingShape(NeuraShape, NeuraShape),
}
pub trait NeuraCombineInputs<T> {
type Combined;
fn combine(&self, inputs: Vec<impl Borrow<T>>) -> Self::Combined;
// TODO:
// fn shape(&self, input_shapes: Vec<NeuraShape>) -> NeuraShape;
}
impl<F: Clone> NeuraCombineInputs<DVector<F>> for NeuraAxisAppend {
type Combined = DVector<F>;
fn combine(&self, inputs: Vec<impl Borrow<DVector<F>>>) -> Self::Combined {
assert!(inputs.len() > 0);
let mut res = Vec::with_capacity(inputs.iter().map(|vec| vec.borrow().len()).sum());
for input in inputs {
for x in input.borrow().iter() {
res.push(x.clone());
}
}
DVector::from_data(VecStorage::new(Dyn(res.len()), Const as Const<1>, res))
}
}
impl NeuraCombineInputs<NeuraShape> for NeuraAxisAppend {
type Combined = Result<NeuraShape, NeuraAxisErr>;
fn combine(&self, inputs: Vec<impl Borrow<NeuraShape>>) -> Self::Combined {
let mut inputs = inputs.into_iter().map(|x| *x.borrow());
if let Some(mut res) = inputs.next() {
for operand in inputs {
match (res, operand) {
(NeuraShape::Vector(x), NeuraShape::Vector(y)) => {
res = NeuraShape::Vector(x + y);
}
(x, y) => {
return Err(NeuraAxisErr::ConflictingShape(x, y));
}
}
}
Ok(res)
} else {
Err(NeuraAxisErr::NoInput)
}
}
}

@ -0,0 +1,153 @@
use super::*;
pub trait NeuraResidualConstruct {
type Constructed;
type Err;
fn construct_residual(
self,
input: NeuraResidualInput<NeuraShape>,
) -> Result<Self::Constructed, Self::Err>;
}
#[derive(Clone, Debug)]
pub enum NeuraResidualConstructErr<LayerErr, ChildErr> {
LayerErr(LayerErr),
ChildErr(ChildErr),
OOBConnection(usize),
AxisErr(NeuraAxisErr),
}
use NeuraResidualConstructErr::*;
impl<Layer: NeuraPartialLayer, Axis> NeuraResidualConstruct for NeuraResidualNode<Layer, (), Axis>
where
Axis: NeuraCombineInputs<NeuraShape, Combined = Result<NeuraShape, NeuraAxisErr>>,
{
type Constructed = NeuraResidualNode<Layer::Constructed, (), Axis>;
type Err = NeuraResidualConstructErr<Layer::Err, ()>;
fn construct_residual(
self,
input: NeuraResidualInput<NeuraShape>,
) -> Result<Self::Constructed, Self::Err> {
let (layer_input_shape, _rest) = input.shift();
let layer_input_shape = self
.axis
.combine(layer_input_shape)
.map_err(|e| AxisErr(e))?;
let layer = self
.layer
.construct(layer_input_shape)
.map_err(|e| LayerErr(e))?;
let layer_shape = layer.output_shape();
if let Some(oob_offset) = self.offsets.iter().copied().find(|o| *o > 0) {
return Err(OOBConnection(oob_offset));
}
// TODO: check rest for non-zero columns
Ok(NeuraResidualNode {
layer,
child_network: (),
offsets: self.offsets,
axis: self.axis,
output_shape: Some(layer_shape),
})
}
}
impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraResidualConstruct, Axis> NeuraResidualConstruct
for NeuraResidualNode<Layer, ChildNetwork, Axis>
where
Axis: NeuraCombineInputs<NeuraShape, Combined = Result<NeuraShape, NeuraAxisErr>>,
{
type Constructed = NeuraResidualNode<Layer::Constructed, ChildNetwork::Constructed, Axis>;
type Err = NeuraResidualConstructErr<Layer::Err, ChildNetwork::Err>;
fn construct_residual(
self,
input: NeuraResidualInput<NeuraShape>,
) -> Result<Self::Constructed, Self::Err> {
let (layer_input_shape, mut rest) = input.shift();
let layer_input_shape = self
.axis
.combine(layer_input_shape)
.map_err(|e| AxisErr(e))?;
let layer = self
.layer
.construct(layer_input_shape)
.map_err(|e| LayerErr(e))?;
let layer_shape = Rc::new(layer.output_shape());
for &offset in &self.offsets {
rest.push(offset, Rc::clone(&layer_shape));
}
let layer_shape = *layer_shape;
let child_network = self
.child_network
.construct_residual(rest)
.map_err(|e| ChildErr(e))?;
Ok(NeuraResidualNode {
layer,
child_network,
offsets: self.offsets,
axis: self.axis,
output_shape: Some(layer_shape),
})
}
}
impl<Layer, Axis> NeuraShapedLayer for NeuraResidualNode<Layer, (), Axis> {
#[inline(always)]
fn output_shape(&self) -> NeuraShape {
self.output_shape.unwrap()
}
}
impl<Layer, ChildNetwork: NeuraShapedLayer, Axis> NeuraShapedLayer
for NeuraResidualNode<Layer, ChildNetwork, Axis>
{
#[inline(always)]
fn output_shape(&self) -> NeuraShape {
self.child_network.output_shape()
}
}
impl<Layers: NeuraShapedLayer> NeuraShapedLayer for NeuraResidual<Layers> {
#[inline(always)]
fn output_shape(&self) -> NeuraShape {
self.layers.output_shape()
}
}
impl<Layers: NeuraResidualConstruct> NeuraPartialLayer for NeuraResidual<Layers>
where
// Should always be satisfied:
Layers::Constructed: NeuraShapedLayer,
{
type Constructed = NeuraResidual<Layers::Constructed>;
type Err = Layers::Err;
fn construct(self, input_shape: NeuraShape) -> Result<Self::Constructed, Self::Err> {
let input_shape = Rc::new(input_shape);
let mut inputs = NeuraResidualInput::new();
for &offset in &self.initial_offsets {
inputs.push(offset, Rc::clone(&input_shape));
}
drop(input_shape);
let layers = self.layers.construct_residual(inputs)?;
Ok(NeuraResidual {
layers,
initial_offsets: self.initial_offsets,
})
}
}

@ -0,0 +1,38 @@
use std::rc::Rc;
#[derive(Clone, Debug)]
pub struct NeuraResidualInput<T> {
// TODO: try to remove this Rc
slots: Vec<Vec<Rc<T>>>,
}
impl<T> NeuraResidualInput<T> {
pub fn new() -> Self {
Self { slots: Vec::new() }
}
pub fn push(&mut self, offset: usize, data: Rc<T>) {
while self.slots.len() <= offset {
self.slots.push(Vec::new());
}
self.slots[offset].push(data);
}
pub fn shift(&self) -> (Vec<Rc<T>>, NeuraResidualInput<T>) {
let res = self.slots.get(0).cloned().unwrap_or(vec![]);
let new_input = Self {
slots: self.slots.iter().skip(1).cloned().collect(),
};
(res, new_input)
}
/// Returns the first input item of the first slot
pub fn get_first(self) -> Option<Rc<T>> {
// TODO: return None if the first slot is bigger than 1 or if there are multiple non-empty slots
self.slots
.into_iter()
.next()
.and_then(|first_slot| first_slot.into_iter().next())
}
}

@ -0,0 +1,259 @@
use std::rc::Rc;
use nalgebra::{DVector, Scalar};
use num::Float;
use crate::layer::*;
mod input;
pub use input::*;
mod axis;
pub use axis::*;
mod construct;
pub use construct::NeuraResidualConstructErr;
#[derive(Clone, Debug, PartialEq)]
pub struct NeuraResidual<Layers> {
/// Instance of NeuraResidualNode
layers: Layers,
/// Array of which layers to send the input to, defaults to `vec![0]`
initial_offsets: Vec<usize>,
}
impl<Layers> NeuraResidual<Layers> {
pub fn new(layers: Layers) -> Self {
Self {
layers,
initial_offsets: vec![0],
}
}
pub fn offset(mut self, offset: usize) -> Self {
self.initial_offsets.push(offset);
self
}
pub fn offsets(mut self, offsets: Vec<usize>) -> Self {
self.initial_offsets = offsets;
self
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct NeuraResidualNode<Layer, ChildNetwork, Axis> {
pub layer: Layer,
pub child_network: ChildNetwork,
/// Array of relative layers indices to send the offset of this layer to,
/// defaults to `vec![0]`.
offsets: Vec<usize>,
pub axis: Axis,
output_shape: Option<NeuraShape>,
}
impl<Layer, ChildNetwork> NeuraResidualNode<Layer, ChildNetwork, NeuraAxisAppend> {
pub fn new(layer: Layer, child_network: ChildNetwork) -> Self {
Self {
layer,
child_network,
offsets: vec![0],
axis: NeuraAxisAppend,
output_shape: None,
}
}
}
impl<Layer, ChildNetwork, Axis> NeuraResidualNode<Layer, ChildNetwork, Axis> {
pub fn offsets(mut self, offsets: Vec<usize>) -> Self {
self.offsets = offsets;
self
}
pub fn offset(mut self, offset: usize) -> Self {
self.offsets.push(offset);
self
}
pub fn axis<Axis2>(self, axis: Axis2) -> NeuraResidualNode<Layer, ChildNetwork, Axis2> {
NeuraResidualNode {
layer: self.layer,
child_network: self.child_network,
offsets: self.offsets,
axis,
// Drop the knowledge of output_shape
output_shape: None,
}
}
}
impl<F: Float + Scalar, Layer, ChildNetwork, Axis> NeuraLayer<NeuraResidualInput<DVector<F>>>
for NeuraResidualNode<Layer, ChildNetwork, Axis>
where
Axis: NeuraCombineInputs<DVector<F>>,
Layer: NeuraLayer<Axis::Combined, Output = DVector<F>>,
ChildNetwork: NeuraLayer<NeuraResidualInput<DVector<F>>>,
{
type Output = <ChildNetwork as NeuraLayer<NeuraResidualInput<DVector<F>>>>::Output;
fn eval(&self, input: &NeuraResidualInput<DVector<F>>) -> Self::Output {
let (inputs, mut rest) = input.shift();
let layer_input = self.axis.combine(inputs);
let layer_output = Rc::new(self.layer.eval(&layer_input));
for &offset in &self.offsets {
rest.push(offset, Rc::clone(&layer_output));
}
self.child_network.eval(&rest)
}
}
impl<F: Clone, Output: Clone, Layers> NeuraLayer<DVector<F>> for NeuraResidual<Layers>
where
Layers: NeuraLayer<NeuraResidualInput<DVector<F>>, Output = NeuraResidualInput<Output>>,
{
type Output = Output;
fn eval(&self, input: &DVector<F>) -> Self::Output {
let input: Rc<DVector<F>> = Rc::new((*input).clone());
let mut inputs = NeuraResidualInput::new();
for &offset in &self.initial_offsets {
inputs.push(offset, Rc::clone(&input));
}
drop(input);
let output = self.layers.eval(&inputs);
let result = output.get_first()
.expect("Invalid NeuraResidual state: network returned no data, did you forget to link the last layer?")
.into();
Rc::unwrap_or_clone(result)
}
}
#[macro_export]
macro_rules! neura_residual {
[ "__combine_layers", ] => {
()
};
[ "__combine_layers",
$layer:expr $(, $axis:expr)? $( => $( $offset:expr ),* )?
$(; $( $rest_layer:expr $(, $rest_axis:expr)? $( => $( $rest_offset:expr ),* )? );*)?
] => {{
let layer = $crate::network::residual::NeuraResidualNode::new($layer,
neura_residual![
"__combine_layers",
$($( $rest_layer $(, $rest_axis)? $( => $( $rest_offset ),* )? );*)?
]
);
$(
let layer = layer.axis($axis);
)?
$(
let layer = layer.offsets(vec![$($offset),*]);
)?
layer
}};
[
$( <= $( $initial_offset:expr ),* ;)?
$( $layer:expr $(, $axis:expr)? $( => $( $offset:expr ),* $(,)? )? );*
$(;)?
] => {{
let res = $crate::network::residual::NeuraResidual::new(
neura_residual![ "__combine_layers", $( $layer $(, $axis)? $( => $( $offset ),* )? );* ]
);
$(
let res = res.offsets(vec![$($initial_offset),*]);
)?
res
}};
}
#[cfg(test)]
mod test {
use nalgebra::dvector;
use crate::neura_layer;
use super::*;
#[test]
fn test_resnet_eval() {
let network = NeuraResidual::new(
NeuraResidualNode::new(
neura_layer!("dense", 4)
.construct(NeuraShape::Vector(2))
.unwrap(),
NeuraResidualNode::new(
neura_layer!("dense", 3)
.construct(NeuraShape::Vector(4))
.unwrap(),
NeuraResidualNode::new(
neura_layer!("dense", 6)
.construct(NeuraShape::Vector(7))
.unwrap(),
(),
),
),
)
.offset(1),
);
network.eval(&dvector![0.2, 0.4]);
}
#[test]
fn test_resnet_macro() {
let network = neura_residual![
<= 0, 2;
neura_layer!("dense", 5) => 0, 1;
neura_layer!("dense", 5);
neura_layer!("dense", 3)
];
println!("{:#?}", network);
assert_eq!(network.initial_offsets, vec![0, 2]);
assert_eq!(network.layers.offsets, vec![0, 1]);
assert_eq!(network.layers.child_network.offsets, vec![0]);
assert_eq!(network.layers.child_network.child_network.child_network, ());
let network = neura_residual![
neura_layer!("dense", 4) => 0;
];
assert_eq!(network.initial_offsets, vec![0]);
}
#[test]
fn test_resnet_partial() {
let network = neura_residual![
<= 0, 1;
neura_layer!("dense", 2) => 0, 1;
neura_layer!("dense", 4);
neura_layer!("dense", 8)
]
.construct(NeuraShape::Vector(1))
.unwrap();
assert_eq!(network.output_shape(), NeuraShape::Vector(8));
network.eval(&dvector![0.0]);
}
}

@ -1,3 +1,5 @@
use crate::layer::NeuraShapedLayer;
use super::*;
pub trait NeuraSequentialConstruct {
@ -40,7 +42,7 @@ impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraSequentialConstruct> NeuraSequ
// TODO: ensure that this operation (and all recursive operations) are directly allocated on the heap
let child_network = self
.child_network
.construct(Layer::output_shape(&layer))
.construct(layer.output_shape())
.map_err(|e| NeuraSequentialConstructErr::Child(e))?;
let child_network = Box::new(child_network);
@ -50,3 +52,19 @@ impl<Layer: NeuraPartialLayer, ChildNetwork: NeuraSequentialConstruct> NeuraSequ
})
}
}
impl<Layer: NeuraShapedLayer> NeuraShapedLayer for NeuraSequential<Layer, ()> {
#[inline(always)]
fn output_shape(&self) -> NeuraShape {
self.layer.output_shape()
}
}
impl<Layer, ChildNetwork: NeuraShapedLayer> NeuraShapedLayer
for NeuraSequential<Layer, ChildNetwork>
{
#[inline(always)]
fn output_shape(&self) -> NeuraShape {
self.child_network.output_shape()
}
}

Loading…
Cancel
Save