general save
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
import os
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
|
||||
import sys
|
||||
|
||||
@@ -199,7 +200,7 @@ def normals_reconst_loss(y_pred, normals_gt, mask):
|
||||
|
||||
|
||||
def train_step(batch):
|
||||
# plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'], batch['joints'], batch['original'], batch['original_mask'])
|
||||
plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'], batch['joints'], batch['original'], batch['original_mask'])
|
||||
# plot3d(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'])
|
||||
### Training
|
||||
model.train() # train mode is on by default after construction
|
||||
@@ -275,8 +276,8 @@ if __name__ == '__main__':
|
||||
batch_size = 1
|
||||
use_cache = True
|
||||
|
||||
train_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
|
||||
valid_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
|
||||
train_dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
|
||||
valid_dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
|
||||
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
|
||||
valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
|
||||
train_accumulator = Accumulator()
|
||||
|
||||
Reference in New Issue
Block a user