Spaces:
Runtime error
Runtime error
File size: 1,834 Bytes
e775f6d |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import torch; torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
import torch.distributions
import matplotlib.pyplot as plt; plt.rcParams['figure.dpi'] = 200
device = 'cuda' if torch.cuda.is_available() else 'cpu'
def get_activation(activation):
if activation == 'tanh':
activ = F.tanh
elif activation == 'relu':
activ = F.relu
elif activation == 'mish':
activ = F.mish
elif activation == 'sigmoid':
activ = torch.sigmoid
elif activation == 'leakyrelu':
activ = F.leaky_relu
elif activation == 'exp':
activ = torch.exp
else:
raise ValueError
return activ
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dims, output_dim, activation, dropout, final_activ=None):
super(SimpleNet, self).__init__()
self.linears = nn.ModuleList()
self.dropouts = nn.ModuleList()
self.output_dim = output_dim
dims = [input_dim] + hidden_dims + [output_dim]
for d_in, d_out in zip(dims[:-1], dims[1:]):
self.linears.append(nn.Linear(d_in, d_out))
self.dropouts.append(nn.Dropout(dropout))
self.activation = get_activation(activation)
self.n_layers = len(self.linears)
self.layer_range = range(self.n_layers)
if final_activ != None:
self.final_activ = get_activation(final_activ)
self.use_final_activ = True
else:
self.use_final_activ = False
def forward(self, x):
for i_layer, layer, dropout in zip(self.layer_range, self.linears, self.dropouts):
x = layer(x)
if i_layer != self.n_layers - 1:
x = self.activation(dropout(x))
if self.use_final_activ: x = self.final_activ(x)
return x
|