48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class DiffusionModel(nn.Module):
|
|
def __init__(self, time_dim: int = 64):
|
|
super(DiffusionModel, self).__init__()
|
|
self.time_dim = time_dim
|
|
self.layers = nn.ModuleList()
|
|
|
|
def pos_encoding(self, t, channels):
|
|
inv_freq = 1.0 / (
|
|
10000 ** (torch.arange(0, channels, 2).float() / channels)
|
|
).to(t.device)
|
|
pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq)
|
|
pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq)
|
|
pos_enc = torch.cat((pos_enc_a, pos_enc_b), dim=-1)
|
|
return pos_enc
|
|
|
|
|
|
def forward(self, x, t, inputs):
|
|
t = t.unsqueeze(-1).type(torch.float)
|
|
t = self.pos_encoding(t, self.time_dim)
|
|
|
|
x = torch.cat((x, t, inputs), dim=-1)
|
|
|
|
for layer in self.layers[:-1]:
|
|
x = layer(x)
|
|
if not isinstance(layer, nn.ReLU):
|
|
x = torch.cat((x, t, inputs), dim=-1)
|
|
|
|
x = self.layers[-1](x)
|
|
return x
|
|
|
|
class SimpleDiffusionModel(DiffusionModel):
|
|
def __init__(self, input_size: int, hidden_sizes: list, other_inputs_dim: int, time_dim: int = 64):
|
|
super(SimpleDiffusionModel, self).__init__(time_dim)
|
|
|
|
self.other_inputs_dim = other_inputs_dim
|
|
|
|
self.layers.append(nn.Linear(input_size + time_dim + other_inputs_dim, hidden_sizes[0]))
|
|
self.layers.append(nn.ReLU())
|
|
|
|
for i in range(1, len(hidden_sizes)):
|
|
self.layers.append(nn.Linear(hidden_sizes[i - 1] + time_dim + other_inputs_dim, hidden_sizes[i]))
|
|
self.layers.append(nn.ReLU())
|
|
|
|
self.layers.append(nn.Linear(hidden_sizes[-1] + time_dim + other_inputs_dim, input_size))
|