78 lines
2.0 KiB
Python
78 lines
2.0 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
|
|
|
|
# Encoder
|
|
self.encoder = nn.Sequential(
|
|
nn.Linear(3*3, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, 256),
|
|
nn.ReLU(),
|
|
nn.Linear(256, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, bottleneck_size),
|
|
nn.ReLU()
|
|
)
|
|
|
|
# Decoder
|
|
self.decoder = nn.Sequential(
|
|
nn.Linear(bottleneck_size, 512),
|
|
nn.ReLU(),
|
|
nn.Linear(512, 256),
|
|
nn.ReLU(),
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 3*3),
|
|
nn.Tanh(),
|
|
)
|
|
|
|
def forward(self, x):
|
|
# Encode
|
|
matrix = self.encoder(x.view(-1, 3*3))
|
|
|
|
# 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 |