hidden_states = []
logging.debug(inputs.size())
for i in range(1, self.blocks + 1):
inputs, state_stage = self.forward_by_stage(
inputs, getattr(self, 'stage' + str(i)),
getattr(self, 'rnn' + str(i)))
hidden_states.append(state_stage)
return tuple(hidden_states)
if __name__ == "__main__":
from net_params import convgru_encoder_params, convgru_decoder_params
from data.mm import MovingMNIST
encoder = Encoder(convgru_encoder_params[0],
convgru_encoder_params[1]).cuda()
trainFolder = MovingMNIST(is_train=True,
root='data/',
n_frames_input=10,
n_frames_output=10,
num_objects=[3])
trainLoader = torch.utils.data.DataLoader(
trainFolder,
batch_size=4,
shuffle=False,
)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for i, (idx, targetVar, inputVar, _, _) in enumerate(trainLoader):
inputs = inputVar.to(device) # B,S,1,64,64
state = encoder(inputs)
random_seed = 1996
np.random.seed(random_seed)
torch.manual_seed(random_seed)
if torch.cuda.device_count() > 1:
torch.cuda.manual_seed_all(random_seed)
else:
torch.cuda.manual_seed(random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
save_dir = "./save_model/" + TIMESTAMP
trainFolder = MovingMNIST(
is_train=True,
root="data/",
n_frames_input=args.frames_input,
n_frames_output=args.frames_output,
num_objects=[3],
)
validFolder = MovingMNIST(
is_train=False,
root="data/",
n_frames_input=args.frames_input,
n_frames_output=args.frames_output,
num_objects=[3],
)
trainLoader = torch.utils.data.DataLoader(trainFolder,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False)
validLoader = torch.utils.data.DataLoader(validFolder,