Spaces:
Sleeping
Sleeping
File size: 1,263 Bytes
3010d5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
'''Boxes for defining and using PyTorch models.'''
import inspect
from . import ops
LAYERS = {}
@ops.op("Define PyTorch model", sub_nodes=LAYERS)
def define_pytorch_model(*, sub_flow):
# import torch # Lazy import because it's slow.
print('sub_flow:', sub_flow)
return 'hello ' + str(sub_flow)
def register_layer(name):
def decorator(func):
sig = inspect.signature(func)
inputs = {
name: param.annotation
for name, param in sig.parameters.items()
if param.kind != param.KEYWORD_ONLY}
params = {
name: param.default if param.default is not inspect._empty else None
for name, param in sig.parameters.items()
if param.kind == param.KEYWORD_ONLY}
outputs = {'x': 'tensor'}
LAYERS[name] = ops.Op(func, name, params=params, inputs=inputs, outputs=outputs, type='vertical')
return func
return decorator
@register_layer('LayerNorm')
def normalization():
return 'LayerNorm'
@register_layer('Dropout')
def dropout(*, p=0.5):
return f'Dropout ({p})'
@register_layer('Linear')
def linear(*, output_dim: int):
return f'Linear {output_dim}'
@register_layer('Graph Convolution')
def graph_convolution():
return 'GraphConv'
@register_layer('Nonlinearity')
def nonlinearity():
return 'ReLU'
|