mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-05-19 08:21:11 +08:00
65 lines
2.4 KiB
Python
65 lines
2.4 KiB
Python
import time
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class GinkaRNN(nn.Module):
|
|
def __init__(self, tile_classes=32, cond_dim=256, input_dim=256, hidden_dim=512, num_layers=1):
|
|
super().__init__()
|
|
|
|
# 输入部分
|
|
self.embedding = nn.Embedding(tile_classes, input_dim)
|
|
self.input_fc = nn.Linear(input_dim, input_dim)
|
|
|
|
self.gru = nn.GRU(input_dim + cond_dim, hidden_dim, num_layers, batch_first=True)
|
|
self.fc = nn.Linear(hidden_dim, tile_classes)
|
|
|
|
def forward(self, x: torch.Tensor, cond: torch.Tensor):
|
|
"""
|
|
x: [B, T]
|
|
cond: [B, cond_dim]
|
|
"""
|
|
B, T = x.shape
|
|
tile_emb = self.input_fc(self.embedding(x)) # [B, T, input_dim]
|
|
cond_expand = cond.unsqueeze(1).expand(B, T, cond.shape[-1]) # [B, T, cond_dim]
|
|
|
|
# 拼接 tile + cond
|
|
step_input = torch.cat([tile_emb, cond_expand], dim=-1)
|
|
|
|
out, _ = self.gru(step_input)
|
|
logits = self.fc(out)
|
|
return logits
|
|
|
|
def print_memory(tag=""):
|
|
print(f"{tag} | 当前显存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB, 最大显存: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
|
|
|
|
if __name__ == "__main__":
|
|
input = torch.rand(1, 32, 32, 32).cuda()
|
|
tag = torch.rand(1, 64).cuda()
|
|
val = torch.rand(1, 16).cuda()
|
|
|
|
# 初始化模型
|
|
model = GinkaRNN().cuda()
|
|
|
|
print_memory("初始化后")
|
|
|
|
# 前向传播
|
|
start = time.perf_counter()
|
|
fake0 = model(input, 0, tag, val)
|
|
fake1 = model(F.softmax(fake0, dim=1), 1, tag, val)
|
|
fake2 = model(F.softmax(fake1, dim=1), 1, tag, val)
|
|
fake3 = model(F.softmax(fake2, dim=1), 1, tag, val)
|
|
end = time.perf_counter()
|
|
|
|
print_memory("前向传播后")
|
|
|
|
print(f"推理耗时: {end - start}")
|
|
print(f"输入形状: feat={input.shape}")
|
|
print(f"输出形状: output={fake3.shape}")
|
|
print(f"Random parameters: {sum(p.numel() for p in model.head.parameters())}")
|
|
print(f"Cond parameters: {sum(p.numel() for p in model.cond.parameters())}")
|
|
print(f"Input parameters: {sum(p.numel() for p in model.input.parameters())}")
|
|
print(f"UNet parameters: {sum(p.numel() for p in model.unet.parameters())}")
|
|
print(f"Output parameters: {sum(p.numel() for p in model.output.parameters())}")
|
|
print(f"Total parameters: {sum(p.numel() for p in model.parameters())}")
|