gcn icon indicating copy to clipboard operation
gcn copied to clipboard

Modularize functions

Open HLWeil opened this issue 3 years ago • 0 comments

Many of the gcn and also the train functions might be generalized to allow for a more modular and broader use.

E.g. the GCNModel.create function might be generalized to a compose function:

Instead of

let create nfeat nhid nclass dropout adj =
    let gc1 = gcnLayer nfeat nhid true adj
    let gc2 = gcnLayer nhid nclass true adj        
    let drp = Dropout(dropout) |> M

    F [gc1;gc2;drp] (fun t ->
        use t = gc1.forward(t)
        use t = Functions.ReLU(t)
        use t = drp.forward(t)
        use t = gc2.forward(t)
        let t = Functions.LogSoftmax(t, dimension=1L)
        t)   

there could be a general layer composition function

let composeWithDropOut (dropOut : #IModel) (activationFunc : TorchTensor -> TorchTensor) (outFunc : TorchTensor -> TorchTensor) (layers : FuncModel list) =
    let rec composeF (f : TorchTensor -> TorchTensor) (remainingLayers : #IModel list) =
        match remainingLayers with
        | [] -> failwith "no layers were given"
        | l :: [] -> f >> l.forward >> outFunc
        | l :: ls ->  composeF (f >> l.forward >> dropOut.forward >> activationFunc) ls
    let forward = composeF id layers
    let models : IModel list =  List.append (layers |> List.map (fun x -> x :> IModel)) [dropOut]
    F models forward

with your specific case resulting in

///Create two layer GCN model with dropout
let create nfeat nhid nclass dropout adj =
    let gc1 = gcnLayer nfeat nhid true adj
    let gc2 = gcnLayer nhid nclass true adj        
    let drp = Dropout(dropout) |> M

    composeWithDropOut drp (fun t -> Functions.ReLU(t)) (fun t -> Functions.LogSoftmax(t,dimension=1L)) [gc1;gc2]

HLWeil avatar Jun 07 '21 17:06 HLWeil