import time
### Hyperparameters ###
seq_len = 10
heatup_seq_len = 10 * 2
batch_size = 4
num_workers = 0
lr = 1e-3 #learning rate
epoch = 10000
displaying = True
weight_path_lstm = "weights_predrnn.chkpt"
### DATALOADER ###
ds = PoolDataset(seq_len=seq_len,
heatup_seq_len=heatup_seq_len,
sum_channels=True,
transform=lambda frames: F.avg_pool2d(
(frames / frames.max() - 0.5) * 2, 4),
variance=10)
dataloader = DataLoader(
ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
#collate_fn=collate_fn,
drop_last=True,
pin_memory=True)
### MODELS ###
lstm = nn.Sequential(
PredRNN(1, 16, 4, 4, num_layers=1),
PredRNN(16, 32, 4, 4, num_layers=1),
from tqdm import tqdm
import os
plt.ion()
### Hyperparameters ###
seq_len = 4
heatup_seq_len = 0
batch_size = 1
num_workers = 8
lr = 1e-3 #learning rate
epoch = 10
displaying = False
weight_path = "weights_e3d_lstm.chkpt"
### DATALOADER ###
ds = PoolDataset(seq_len=seq_len, heatup_seq_len=heatup_seq_len)
dataloader = DataLoader(
ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
#collate_fn=collate_fn,
drop_last=True)
### MODELS ###
e3d_lstm = E3DLSTM((1, seq_len, 256, 256), 16, 1, (1, 3, 3), 2)
if os.path.isfile(weight_path):
w = torch.load(weight_path)
e3d_lstm.load_state_dict(w['e3d_lstm'])
del w
e3d_lstm = e3d_lstm
import time
plt.ion()
### Hyperparameters ###
seq_len = 100
heatup_seq_len = 0
batch_size = 1
num_workers = 8 * 2
lr = 1e-2 #learning rate
epoch = 1
displaying = True
weight_path_conv3d = "weights_conv3d.chkpt"
### DATALOADER ###
ds = PoolDataset(seq_len=seq_len,
heatup_seq_len=heatup_seq_len,
sum_channels=False)
dataloader = DataLoader(
ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
#collate_fn=collate_fn,
drop_last=True)
### MODELS ###
model = Conv3D()
if os.path.isfile(weight_path_conv3d):
w = torch.load(weight_path_conv3d)
model.load_state_dict(w['model'])
import os
import time
### Hyperparameters ###
seq_len = 4
heatup_seq_len = 10
batch_size = 1
num_workers = 8 * 2
lr = 1e-3 #learning rate
epoch = 10
displaying = True
weight_path_lstm = "weights_convlstm.chkpt"
### DATALOADER ###
ds = PoolDataset(seq_len=seq_len,
heatup_seq_len=heatup_seq_len,
sum_channels=True,
transform=lambda frames: (frames / frames.max() - 0.5) * 2)
dataloader = DataLoader(
ds,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
#collate_fn=collate_fn,
drop_last=True,
pin_memory=True)
### MODELS ###
lstm = nn.Sequential(ConvLSTM(1, 32, 5), ConvLSTM(32, 1, 1))
if os.path.isfile(weight_path_lstm):
w = torch.load(weight_path_lstm)