Analysis Software
Documentation for
sPHENIX
simulation software
Home page
Related Pages
Modules
Namespaces
Classes
Files
Examples
External Links
File List
File Members
All
Classes
Namespaces
Files
Functions
Variables
Typedefs
Enumerations
Enumerator
Friends
Macros
Groups
Pages
TorchEdgeClassifier.hpp
Go to the documentation of this file.
Or view
the newest version in sPHENIX GitHub for file TorchEdgeClassifier.hpp
1
// This file is part of the Acts project.
2
//
3
// Copyright (C) 2023 CERN for the benefit of the Acts project
4
//
5
// This Source Code Form is subject to the terms of the Mozilla Public
6
// License, v. 2.0. If a copy of the MPL was not distributed with this
7
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
8
9
#pragma once
10
11
#include "
Acts/Plugins/ExaTrkX/Stages.hpp
"
12
#include "
Acts/Utilities/Logger.hpp
"
13
14
#include <memory>
15
16
namespace
torch::jit {
17
class
Module;
18
}
19
20
namespace
c10 {
21
enum class
DeviceType : int8_t;
22
}
23
24
namespace Acts {
25
26
class
TorchEdgeClassifier
final :
public
Acts::EdgeClassificationBase
{
27
public
:
28
struct
Config
{
29
std::string
modelPath
;
30
int
numFeatures = 3;
31
float
cut = 0.21;
32
int
nChunks = 1;
// NOTE for GNN use 1
33
bool
undirected =
false
;
34
};
35
36
TorchEdgeClassifier
(
const
Config
&
cfg
, std::unique_ptr<const Logger>
logger
);
37
~
TorchEdgeClassifier
();
38
39
std::tuple<std::any, std::any, std::any> operator()(
40
std::any nodes, std::any
edges
,
int
deviceHint = -1)
override
;
41
42
Config
config
()
const
{
return
m_cfg
; }
43
44
private
:
45
std::unique_ptr<const Acts::Logger>
m_logger
;
46
const
auto
&
logger
()
const
{
return
*
m_logger
; }
47
48
Config
m_cfg
;
49
c10::DeviceType
m_deviceType
;
50
std::unique_ptr<torch::jit::Module>
m_model
;
51
};
52
53
}
// namespace Acts
acts
blob
sPHENIX
Plugins
ExaTrkX
include
Acts
Plugins
ExaTrkX
TorchEdgeClassifier.hpp
Built by
Jin Huang
. updated:
Sat Feb 17 2024 22:17:41
using
1.8.2 with
sPHENIX GitHub integration