general save

This commit is contained in:
2023-05-15 17:14:49 +00:00
parent a8647951aa
commit 178fa676ca
14 changed files with 83507 additions and 50 deletions
+5 -4
View File
@@ -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()