[demos] Simplify network_training_pytorch.py
@@ -6,32 +6,14 @@ from phi.torch.flow import *
# TORCH_BACKEND.set_default_device('GPU')
-
-
-class Net(torch.nn.Module):
-
- def __init__(self):
- super(Net, self).__init__()
- self.conv1 = torch.nn.Conv2d(2, 4, 3, padding=1)
- self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1)
- self.conv3 = torch.nn.Conv2d(4, 2, 3, padding=1)
-
- def forward(self, x):
- x = torchf.relu(self.conv1(x))
- x = torchf.relu(self.conv2(x))
- x = self.conv3(x)
- return x
-
-
-net = Net().to(TORCH_BACKEND.get_default_device().ref)
+net = u_net(2, 2)
optimizer = optim.Adam(net.parameters(), lr=1e-3)
-# Initialize variables for ModuleViewer
DOMAIN = Domain(x=64, y=64)
prediction = DOMAIN.vector_grid(0)
prediction_div = DOMAIN.scalar_grid(0)
-
app = ModuleViewer()
+
for step in app.range(100, warmup=1):
# Load or generate training data
data = DOMAIN.vector_grid(Noise(batch=8, vector=2))
