def __init__(self,
in_channels,
out_channels,
blocks=0,
unroll=0,
residual=True,
batch_norm=True,
transfer_data=True):
super(UpConv, self).__init__()
self.residual = residual
self.bn = []
self.unroll = None
self.transfer_data = transfer_data
self.up_conv = MeshConv(in_channels, out_channels)
if transfer_data:
self.conv1 = MeshConv(2 * out_channels, out_channels)
else:
self.conv1 = MeshConv(out_channels, out_channels)
self.conv2 = []
for _ in range(blocks):
self.conv2.append(MeshConv(out_channels, out_channels))
self.conv2 = nn.ModuleList(self.conv2)
if batch_norm:
for _ in range(blocks + 1):
self.bn.append(nn.InstanceNorm2d(out_channels))
self.bn = nn.ModuleList(self.bn)
if unroll:
self.unroll = MeshUnpool(unroll)
def __init__(self,
in_channels,
out_channels,
blocks=0,
unroll=0,
residual=True,
batch_norm=True,
nv=600):
super(UpConv2, self).__init__()
self.residual = residual
self.bn = []
self.unroll = None
self.up_conv = MeshConv(in_channels, out_channels)
self.blocks = blocks
self.conv1 = []
for _ in range(blocks):
self.conv1.append(MeshConv(out_channels, out_channels))
self.conv1 = nn.ModuleList(self.conv1)
self.conv3 = MeshConv(out_channels, 3)
self.conv2 = []
for _ in range(blocks):
self.conv2.append(MeshConv(out_channels, out_channels))
self.conv2 = nn.ModuleList(self.conv2)
if batch_norm:
for _ in range(blocks + 1):
self.bn.append(nn.InstanceNorm2d(out_channels))
self.bn.append(nn.InstanceNorm2d(3))
self.bn = nn.ModuleList(self.bn)
if unroll:
self.unroll = MeshUnpoolLearned(unroll, nv)
self.optimisor = MeshUnpoolOptimisor(name=str(out_channels))
self.flip_edge = MeshUnpoolValenceOptimisor(conv=True,
channel=out_channels)
def __init__(self, in_channels, out_channels, skips=1):
super(MResConv, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.skips = skips
self.conv0 = MeshConv(self.in_channels, self.out_channels, bias=False)
for i in range(self.skips):
setattr(self, 'bn{}'.format(i + 1), nn.BatchNorm2d(self.out_channels))
setattr(self, 'conv{}'.format(i + 1),
MeshConv(self.out_channels, self.out_channels, bias=False))
def __init__(self, in_channels, out_channels, blocks=0, pool=0):
super(DownConv, self).__init__()
self.bn = []
self.pool = None
self.conv1 = MeshConv(in_channels, out_channels)
self.conv2 = []
for _ in range(blocks):
self.conv2.append(MeshConv(out_channels, out_channels))
self.conv2 = nn.ModuleList(self.conv2)
for _ in range(blocks + 1):
self.bn.append(nn.InstanceNorm2d(out_channels))
self.bn = nn.ModuleList(self.bn)
if pool:
self.pool = MeshPool(pool)
def __init__(self,
embd_size=16,
sa_window_size=10): # TODO: sa_window_size
super(MeshTransformerNet, self).__init__()
self.k = [5, 10, 20, 40]
self.res = [600, 450, 300, 210]
self.edges_embedding = MeshEdgeEmbeddingLayer(self.k[0], embd_size)
self.k[0] = embd_size
self.sa_layer = MeshSA(self.k[0],
self.k[0],
window_size=sa_window_size)
self.dropout = nn.Dropout(p=0.2)
for i, ki in enumerate(self.k[:-1]):
# setattr(self, 'sa{}'.format(i), MeshSA(self.k[i + 1], self.k[i + 1], window_size=sa_window_size))
setattr(self, 'sa{}'.format(i),
MeshSA(ki, ki, window_size=sa_window_size))
# setattr(self, 'cirLSTM{}'.format(i), CircularMeshLSTM(ki, 220, self.k[i+1], 1))
# setattr(self, 'conv{}'.format(i), MeshConv(ki, self.k[i + 1]))
setattr(self, 'conv{}'.format(i), MeshConv(ki, self.k[i + 1]))
# setattr(self, 'pool{}'.format(i), MeshPool(self.res[i]))
setattr(self, 'sa_pool{}'.format(i), MeshPoolSA(self.res[i + 1]))
self.gp = torch.nn.AvgPool1d(self.res[-1])
# self.gp = torch.nn.MaxPool1d(self.res[-1])
self.fc = nn.Linear(self.k[-1], 35)
self.relu = nn.ReLU()
def __init__(self,
n_edges,
in_ch=6,
convs=[32, 64],
pool=[],
res_blocks=0,
init_verts=None,
transfer_data=False,
leaky=0,
init_weights_size=0.002):
super(PriorNet, self).__init__()
# check that the number of pools and convs match such that there is a pool between each conv
down_convs = [in_ch] + convs
up_convs = convs[::-1] + [in_ch]
pool_res = [n_edges] + pool
self.encoder_decoder = MeshEncoderDecoder(pools=pool_res,
down_convs=down_convs,
up_convs=up_convs,
blocks=res_blocks,
transfer_data=transfer_data,
leaky=leaky)
self.last_conv = MeshConv(6, 6)
init_weights(self, 'normal', init_weights_size)
eps = 1e-8
self.last_conv.conv.weight.data.uniform_(-1 * eps, eps)
self.last_conv.conv.bias.data.uniform_(-1 * eps, eps)
self.init_verts = init_verts
def __init__(self, in_channels, out_channels, blocks=0, pool=0, nl_block = 0):
super(DownConv, self).__init__()
self.bn = []
self.pool = None
self.nl_block = None
self.conv1 = MeshConv(in_channels, out_channels)
self.conv2 = []
for _ in range(blocks):
self.conv2.append(MeshConv(out_channels, out_channels))
self.conv2 = nn.ModuleList(self.conv2)
for _ in range(blocks + 1):
self.bn.append(nn.InstanceNorm2d(out_channels))
self.bn = nn.ModuleList(self.bn)
if pool:
self.pool = MeshPool(pool)
if nl_block:
# Add non-local block before last downpool
self.nl_block = NLBlock(out_channels, block_type = nl_block)
def __init__(self,
in_channels,
out_channels,
blocks=0,
pool=0,
window_size=50,
num_heads=8):
super(DownConv, self).__init__()
self.bn = []
self.pool = None
self.conv1 = MeshConv(in_channels, out_channels)
self.conv2 = []
for _ in range(blocks):
self.conv2.append(MeshConv(out_channels, out_channels))
self.conv2 = nn.ModuleList(self.conv2)
for _ in range(blocks + 1):
self.bn.append(nn.InstanceNorm2d(out_channels))
self.bn = nn.ModuleList(self.bn)
if pool:
self.self_attention = MeshSelfAttention(out_channels,
64,
window_size=window_size,
heads=num_heads)
self.pool = MeshPoolSA(pool)
parser.add_argument('--norm')
parser.add_argument('--resblocks')
opt = parser.parse_args()
mesh_path = 'tree_98.obj'
mesh = mesh.Mesh(file=mesh_path, opt=opt, export_folder='.')
print(mesh.edges_count)
import torch
from models.layers.mesh_conv import MeshConv, MeshEdgeEmbeddingLayer
from models.layers.mesh_self_attention import MeshSelfAttention
# from models.layers.mesh_conv_old import MeshConv
sa = MeshSelfAttention(5, 5, 10)
mc = MeshConv(5, 15)
# a = torch.rand(1, 750, 5)
import pickle
with open('input_a.p', 'rb') as f:
a = torch.load(f, map_location='cpu')
embd_layer = MeshEdgeEmbeddingLayer(5, 10)
print(a['x'].shape)
print(embd_layer(a['x']).shape)
print(sa(a['x'])[0].shape)
print(a['x'].shape)
o = mc(a['x'], a['mesh'])
print(o.shape)
from torch import nn
def __init__(self, in_feat, out_feat, k=1):
super(ConvBlock, self).__init__()
self.lst = [MeshConv(in_feat, out_feat)]
for i in range(k - 1):
self.lst.append(MeshConv(out_feat, out_feat))
self.lst = nn.ModuleList(self.lst)
