Commit 211679fd by xlwang

correct data loader error

parent f1d77da5
...@@ -22,7 +22,7 @@ def train(model, train_loader, epoch, optimizer, criterion, clip): ...@@ -22,7 +22,7 @@ def train(model, train_loader, epoch, optimizer, criterion, clip):
epoch_loss = 0 epoch_loss = 0
cnt = 0 cnt = 0
loop = tqdm(enumerate(train_loader), total=num_train_iteration_per_epoch) loop = tqdm(enumerate(train_loader.get_iterator()), total=num_train_iteration_per_epoch)
# for _, (x, y) in enumerate(train_loader): # for _, (x, y) in enumerate(train_loader):
for _, (x, y) in loop: for _, (x, y) in loop:
# x/y shape (50, 12, 207, 2) # x/y shape (50, 12, 207, 2)
...@@ -48,7 +48,7 @@ def evaluate(model, val_loader, epoch, criterion): ...@@ -48,7 +48,7 @@ def evaluate(model, val_loader, epoch, criterion):
model.eval() model.eval()
epoch_loss = 0 epoch_loss = 0
cnt = 0 cnt = 0
loop = tqdm(enumerate(val_loader), total=num_val_iteration_per_epoch) loop = tqdm(enumerate(val_loader.get_iterator()), total=num_val_iteration_per_epoch)
with torch.no_grad(): with torch.no_grad():
for i, (x, y) in loop: for i, (x, y) in loop:
cnt += 1 cnt += 1
...@@ -74,7 +74,7 @@ def test(model, test_loader, scaler): ...@@ -74,7 +74,7 @@ def test(model, test_loader, scaler):
predictions = [] predictions = []
with torch.no_grad(): with torch.no_grad():
for i, (x, y) in enumerate(test_loader): for i, (x, y) in enumerate(test_loader.get_iterator()):
x = torch.FloatTensor(x).cuda() x = torch.FloatTensor(x).cuda()
y = torch.FloatTensor(y).cuda() y = torch.FloatTensor(y).cuda()
outputs = model(x, y, 0) # (seq_length+1, batch_size, num_nodes*output_dim) (13, 50, 207*1) outputs = model(x, y, 0) # (seq_length+1, batch_size, num_nodes*output_dim) (13, 50, 207*1)
...@@ -144,9 +144,9 @@ if __name__ == '__main__': ...@@ -144,9 +144,9 @@ if __name__ == '__main__':
scaler = data['scaler'] scaler = data['scaler']
train_data_loader = data['train_loader'].get_iterator() train_data_loader = data['train_loader']
val_data_loader = data['val_loader'].get_iterator() val_data_loader = data['val_loader']
test_data_loader = data['test_loader'].get_iterator() test_data_loader = data['test_loader']
# Initialize model # Initialize model
model = DCGRUModel(batch_size=batch_size, enc_input_dim=input_dim, dec_input_dim=output_dim, model = DCGRUModel(batch_size=batch_size, enc_input_dim=input_dim, dec_input_dim=output_dim,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment