File size: 511 Bytes
3978e51 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from torch import nn
import torch
class FiLM(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, gamma, beta):
return gamma * x + beta
class BTFBroadcastedFiLM(nn.Module):
def __init__(self):
super().__init__()
self.film = FiLM()
def forward(self, x, gamma, beta):
gamma = gamma[None, None, None, :]
beta = beta[None, None, None, :]
return self.film(x, gamma, beta)
|