Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update find and replace to only use nodeids #86

Closed
wants to merge 4 commits into from
Closed

Update find and replace to only use nodeids #86

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cf6eee6
Comment broken parts
BradenEverson Aug 5, 2024
4c50d99
Merge now replaces param nodeid with another nodeid, no strings attac…
BradenEverson Aug 6, 2024
b1d3e9d
Got one replacement working I think
BradenEverson Aug 6, 2024
ff733b6
Replace other find and replace usage
BradenEverson Aug 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/graph/consteval.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,7 @@ impl Context {
}
}
}
_ => {todo!("Match sin and cos")}
}
}

Expand Down Expand Up @@ -724,6 +725,7 @@ impl Context {
}
}
Operation::Constant(_) | Operation::Parameter(_) => {}
_ => {todo!("Match sin and cos")}
}
visitied.insert(node_id);
}
Expand Down
180 changes: 80 additions & 100 deletions src/graph/merge.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ impl Context {
Operation::Log(a) => {
self.nodes[new_id].operation = Operation::Log(old_to_new[&a])
}
Operation::Cos(a) => {
self.nodes[new_id].operation = Operation::Cos(old_to_new[&a])
}
Operation::Sin(a) => {
self.nodes[new_id].operation = Operation::Sin(old_to_new[&a])
}
Operation::ZerosLike(a) => {
self.nodes[new_id].operation = Operation::ZerosLike(old_to_new[&a])
}
Expand Down Expand Up @@ -161,114 +167,88 @@ impl Context {

pub fn find_and_replace_params(
&mut self,
param_reps: &[(&str, &[NodeIdentifier])],
) -> Result<()> {
for (param_name, rep_with) in param_reps {
let params_with_name: Vec<NodeIdentifier> = self
.nodes
.clone()
.into_iter()
.filter(|(_, node)| match node.operation.clone() {
Operation::Parameter(name) => name.contains(param_name),
_ => false,
})
.map(|(id, _)| id)
.collect();
param_reps: &[(NodeIdentifier, NodeIdentifier)],
) -> Result<()> {
for (param, rep_with) in param_reps {
let param_node = &self.nodes[*param];
let rep_with_node = &self.nodes[*rep_with];

if params_with_name.len() != rep_with.len() {
return Err(super::ContextError::IncorrectOutputSizeError(
rep_with.len(),
params_with_name.len(),
));
if param_node.shape != rep_with_node.shape || param_node.dtype != rep_with_node.dtype {
return Err(super::ContextError::InvalidFuseTargetsError(param_node.dtype, rep_with_node.dtype))
}

for i in 0..params_with_name.len() {
let param_node = &self.nodes[params_with_name[i]];
let rep_with_node = &self.nodes[rep_with[i]];

if param_node.shape != rep_with_node.shape || param_node.dtype != rep_with_node.dtype {
return Err(super::ContextError::InvalidFuseTargetsError(param_node.dtype, rep_with_node.dtype))
}
self.nodes[params_with_name[i]] = self.nodes[rep_with[i]].clone();

let param_idx = self.parameters.iter().enumerate().find(|(_, node)| node == &&params_with_name[i]);
if let Some((id, _)) = param_idx {
self.parameters.remove(id);
}
let param_idx = self.parameters.iter().enumerate().find(|(_, node)| node == &param);

/*let node_ext = self.dependent_nodes.get(&rep_with[i]).unwrap_or(&vec![]).clone();
if let Some(node_deps) = self.dependent_nodes.get_mut(&params_with_name[i]) {
node_deps.extend(node_ext.iter())
}*/
if let Some((id, _)) = param_idx {
self.parameters.remove(id);
}

//Add param nodeid to dependent nodes of new node's operation
match self.nodes[params_with_name[i]].operation.clone() {
Operation::Add(a, b)
| Operation::Pow(a, b)
| Operation::Sub(a, b)
| Operation::Mul(a, b)
| Operation::MatMul(a, b)
| Operation::Div(a, b)
| Operation::GreaterThanEq(a, b)
| Operation::GreaterThan(a, b)
| Operation::Equal(a, b)
| Operation::NotEqual(a, b)
| Operation::LessThan(a, b)
| Operation::LessThanEq(a, b)
| Operation::RngUniform(a, b, _)
| Operation::RngNormal(a, b, _) => {
self.dependent_nodes.entry(a).or_insert_with(Vec::new).push(params_with_name[i]);
self.dependent_nodes.entry(b).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::StopGradient(a)
| Operation::Neg(a)
| Operation::Log(a)
| Operation::Exp(a)
| Operation::ZerosLike(a)
| Operation::OneHot(a)
| Operation::TypeCast(a, _)
| Operation::Reshape(a)
| Operation::Transpose(a, _) => {
self.dependent_nodes.entry(a).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::Select {
pred,
on_false,
on_true,
} => {
self.dependent_nodes.entry(pred).or_insert_with(Vec::new).push(params_with_name[i]);
self.dependent_nodes.entry(on_true).or_insert_with(Vec::new).push(params_with_name[i]);
self.dependent_nodes.entry(on_false).or_insert_with(Vec::new).push(params_with_name[i]);
match param_node.operation.clone() {
Operation::Add(a, b)
| Operation::Pow(a, b)
| Operation::Sub(a, b)
| Operation::Mul(a, b)
| Operation::MatMul(a, b)
| Operation::Div(a, b)
| Operation::GreaterThanEq(a, b)
| Operation::GreaterThan(a, b)
| Operation::Equal(a, b)
| Operation::NotEqual(a, b)
| Operation::LessThan(a, b)
| Operation::LessThanEq(a, b)
| Operation::RngUniform(a, b, _)
| Operation::RngNormal(a, b, _) => {
self.dependent_nodes.entry(a).or_insert_with(Vec::new).push(*rep_with);
self.dependent_nodes.entry(b).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceMax { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
Operation::StopGradient(a)
| Operation::Neg(a)
| Operation::Log(a)
| Operation::Exp(a)
| Operation::ZerosLike(a)
| Operation::OneHot(a)
| Operation::TypeCast(a, _)
| Operation::Reshape(a)
| Operation::Transpose(a, _) => {
self.dependent_nodes.entry(a).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceArgmax {
node,
dim: _,
} => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::ReduceSum { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::ReduceMean { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::SliceInDim {
node,
start: _,
stop: _,
stride: _,
dim: _,
} => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
}
Operation::TileInDim { node, n_tiles: _, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(params_with_name[i]);
}
_ => {} //Constants, parameters, don't need nodeid replacement
Operation::Select {
pred,
on_false,
on_true,
} => {
self.dependent_nodes.entry(pred).or_insert_with(Vec::new).push(*rep_with);
self.dependent_nodes.entry(on_true).or_insert_with(Vec::new).push(*rep_with);
self.dependent_nodes.entry(on_false).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceMax { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceArgmax {
node,
dim: _,
} => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceSum { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
Operation::ReduceMean { node, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
Operation::SliceInDim {
node,
start: _,
stop: _,
stride: _,
dim: _,
} => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
Operation::TileInDim { node, n_tiles: _, dim: _ } => {
self.dependent_nodes.entry(node).or_insert_with(Vec::new).push(*rep_with);
}
_ => {} //Constants, parameters, don't need nodeid replacement
}
}

Expand Down
3 changes: 2 additions & 1 deletion src/graph/subterm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,8 @@ impl Context {
to_visit.push(alpha);
to_visit.push(beta);
to_visit.push(x);
}
},
_ => {todo!("Match sin and cos")}
}
node_map.insert(self.nodes[node_id].clone(), node_id);
}
Expand Down
22 changes: 20 additions & 2 deletions src/models/supervised.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ pub struct SupervisedModel {
// separate context which takes parameters, outputs, and targets
pub(crate) compute_metrics: Context,
pub(crate) metric_names: Vec<String>,
pub(crate) metric_inputs: Vec<NodeIdentifier>,
// additional inputs to compute_metrics as the targets of the supervised learning algorithm
pub(crate) targets: Vec<NodeIdentifier>,
// index into compute_metrics context to find differentiable loss function
Expand All @@ -50,6 +51,7 @@ impl SupervisedModel {
inputs: Vec<NodeIdentifier>,
outputs: Vec<NodeIdentifier>,
compute_metrics: Context,
metric_inputs: Vec<NodeIdentifier>,
metric_names: Vec<String>,
targets: Vec<NodeIdentifier>,
loss: NodeIdentifier,
Expand All @@ -68,8 +70,23 @@ impl SupervisedModel {
//compute_metrics will take in outputs and targets as inputs
//outputs is a direct output of inference context
//targets are supplied in constructor
let loss_update = eval_context.merge_graphs(&compute_metrics, &[loss])?[0];
eval_context.find_and_replace_params(&[("outputs", &outputs), ("targets", &targets)])?;
let mut outputs_and_targets_orig_network = outputs.clone();
outputs_and_targets_orig_network.extend(targets.iter());

if outputs_and_targets_orig_network.len() != metric_inputs.len() {
todo!("Better error handling here")
}

let mut desired_new_nodeids = vec![loss];
desired_new_nodeids.extend(metric_inputs.iter());

let new_compute_metric_nodes = eval_context.merge_graphs(&compute_metrics, &desired_new_nodeids)?;
let loss_update = new_compute_metric_nodes[0];

let new_metric_inputs = new_compute_metric_nodes[1..].to_vec();
let fused_replacements: Vec<(NodeIdentifier, NodeIdentifier)> = new_metric_inputs.into_iter().zip(outputs_and_targets_orig_network).collect();

eval_context.find_and_replace_params(&fused_replacements)?;

let evaluation_computation =
eval_context.build("evaluation_computation", vec![loss_update])?;
Expand All @@ -92,6 +109,7 @@ impl SupervisedModel {
inputs,
outputs,
compute_metrics,
metric_inputs,
metric_names,
targets,
loss: loss_update,
Expand Down
3 changes: 2 additions & 1 deletion src/models/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,7 @@ mod tests {
println!("{:?}", rust_result);
}

/*
#[test]
fn test_param_replace() {
let mut f = Context::new();
Expand Down Expand Up @@ -248,6 +249,6 @@ mod tests {
let rust_result = untupled_result.to_vec::<f32>().expect("to_vec");

assert_eq!(16f32, rust_result[0])
}
}*/

}
24 changes: 15 additions & 9 deletions src/training/optimizers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,8 @@ enum LRSchedOffset {
impl LearningRateSchedule {
fn offset(&self) -> LRSchedOffset {
fn zero_offset(sched: &LearningRateSchedule) -> LRSchedOffset {
match sched {
todo!("Can't move out of sched")
/*match sched {
&LearningRateSchedule::Constant(lr) => LRSchedOffset::Constant(lr, 0),
&LearningRateSchedule::ExpDecay {
init_lr,
Expand Down Expand Up @@ -98,11 +99,12 @@ impl LearningRateSchedule {
offset: 0,
}
}
}
}*/
}

fn recurse(sched: &LRSchedOffset, off: usize) -> LRSchedOffset {
match sched {
todo!("Cannot move out of sched")
/*match sched {
&LRSchedOffset::Constant(lr, offset) => LRSchedOffset::Constant(lr, off + offset),
&LRSchedOffset::ExpDecay {
init_lr,
Expand Down Expand Up @@ -141,7 +143,7 @@ impl LearningRateSchedule {
offset: off + offset,
}
}
}
}*/
}
recurse(&zero_offset(self), 0)
}
Expand All @@ -152,7 +154,8 @@ impl LRSchedOffset {
let mut scheduler = Context::new();
let iteration = scheduler.parameter("iteration", [], ElementType::U32)?;

match self {
todo!("Cannot move out of self")
/*match self {
&Self::Constant(lr, _) => {
let lr_node = scheduler.scalar(lr, ElementType::F32)?;
Ok((iteration, scheduler, lr_node))
Expand Down Expand Up @@ -207,7 +210,7 @@ impl LRSchedOffset {
let final_lr = context_1.select(pred, lr_2, lr_1)?;
Ok((inp_1, context_1, final_lr))
}
}
}*/
}
}

Expand Down Expand Up @@ -282,13 +285,16 @@ impl Optimizer<LearningRateSchedule> for SGD {
&self.new_params
}
fn get_old_state(&self) -> &Vec<NodeIdentifier> {
&vec![self.old_iter]
todo!("Cannot reference temporary value")
//&vec![self.old_iter]
}
fn get_new_state(&self) -> &Vec<NodeIdentifier> {
&vec![self.new_iter]
todo!("Cannot reference temporary value")
//&vec![self.new_iter]
}
fn get_user_params(&self) -> LearningRateSchedule {
self.lr_schedule
todo!("Cannot move out of self")
//self.lr_schedule
}
// TODO WILL FAIL NEED PROPER GRAPH MERGING
fn new(lr_schedule: LearningRateSchedule, model_params: Vec<NodeIdentifier>, model: &Context) -> SGD {
Expand Down
Loading
Loading