108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
import os
|
|
import sys
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
|
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
|
|
|
import torch
|
|
import numpy as np
|
|
from torch import nn, Tensor
|
|
from torch.nn import functional as F
|
|
from functools import partial
|
|
from typing import Any, Callable, List, Optional, Sequence
|
|
|
|
torch.set_float32_matmul_precision('high')
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
|
|
|
|
class ViewEncoder(nn.Module):
|
|
def __init__(self, bottleneck_size):
|
|
super(ViewEncoder, self).__init__()
|
|
self.bottleneck_size = bottleneck_size
|
|
|
|
self.encoder = nn.ModuleList([
|
|
nn.Sequential(
|
|
nn.Linear(3*3, 64),
|
|
nn.BatchNorm1d(64),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(64, 128),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(128, 256),
|
|
nn.BatchNorm1d(256),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(256, 512),
|
|
nn.BatchNorm1d(512),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(512, bottleneck_size),
|
|
nn.BatchNorm1d(bottleneck_size),)
|
|
])
|
|
|
|
# Decoder
|
|
self.decoder = nn.ModuleList([
|
|
nn.Sequential(
|
|
nn.Linear(bottleneck_size, 512),
|
|
nn.BatchNorm1d(512),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(512*2, 256),
|
|
nn.BatchNorm1d(256),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(256*2, 128),
|
|
nn.BatchNorm1d(128),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(128*2, 64),
|
|
nn.BatchNorm1d(64),
|
|
nn.ReLU(),),
|
|
nn.Sequential(
|
|
nn.Linear(64*2, 3*3),
|
|
nn.Tanh(),),
|
|
])
|
|
|
|
def forward(self, x):
|
|
# Encode
|
|
f1 = self.encoder[0](x.view(-1, 3*3))
|
|
f2 = self.encoder[1](f1)
|
|
f3 = self.encoder[2](f2)
|
|
f4 = self.encoder[3](f3)
|
|
matrix = self.encoder[4](f4)
|
|
|
|
d1 = self.decoder[0](matrix)
|
|
d2 = self.decoder[1](torch.cat((d1, f4), dim=1))
|
|
d3 = self.decoder[2](torch.cat((d2, f3), dim=1))
|
|
d4 = self.decoder[3](torch.cat((d3, f2), dim=1))
|
|
x = self.decoder[4](torch.cat((d4, f1), dim=1))
|
|
|
|
x = x * np.pi
|
|
return x.view(-1, 3, 3), matrix
|
|
|
|
# # Decode
|
|
# x = self.decoder(matrix)
|
|
# x = x * np.pi
|
|
# return x.view(-1, 3, 3), matrix
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# test the model
|
|
model = ViewEncoder(bottleneck_size=128).to('cuda')
|
|
# model = torch.compile(model)
|
|
print(model)
|
|
print(sum(p.numel() for p in model.parameters() if p.requires_grad))
|
|
print(sum(p.numel() for p in model.parameters()))
|
|
|
|
x = torch.rand((4, 3, 3)).to('cuda')
|
|
output = model(x)
|
|
|
|
# from torchviz import make_dot
|
|
# graph = make_dot(output, params=dict(model.named_parameters()))
|
|
# graph.render("model_graph", format="pdf")
|
|
|
|
pass |