# Hyper-parameters
num_epochs = 5000
lr_sch = 5
learning_rate = 1e-3
batch_size = 16
save_path = f"./langan_20241129_wobg250"

use_tensorboard = True
use_ddp = True
use_longtail = False
use_distillation = False
use_multiscale = True
use_randomcrop = True
# step -1: Initial DDP
if use_ddp:
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    device_id = rank % torch.cuda.device_count()
    device = torch.device(f'cuda:{device_id}')
if os.path.exists(save_path) or device_id != 0:
    pass
else:
    os.makedirs(save_path, exist_ok= True)
print(f"Start running basic DDP example on rank {rank}.")
# step1: Data Loading
# txtname = "distillation_list.txt"
# txtname = "/home/notebook/data/group/zgf/code/data_process/data_list/all/data_list_ip14pc_zed_0517.txt"
# 放大数据集
# txtname = "langan_all_data_sample2.txt"
txtname = "/home/notebook/data/group/Ruby/depthrefine/langanero_files.txt"
data_using = MonoDepthDataset(txtname = txtname)
# print(len(data_using))
if use_ddp:
    sampler = DistributedSampler(data_using, drop_last=True)
    train_loader = DataLoader(data_using, batch_size=batch_size, shuffle=(sampler is None), drop_last=True, sampler=sampler, num_workers=0 )
else:
    train_loader = DataLoader(data_using, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=0)

# model=Mono_Depth().to(device)
model=build_efficientnet_lite('efficientnet_lite3', 1000).to(device)
ckpt = torch.load('/home/notebook/data/group/Ruby/depthrefine/langan_20241127_langan0.3_woprt/1183-6.6036.pth')

oldCkpt = ckpt
newCkpt = {}
pretrained_state_dict =oldCkpt
model_state_dict=model.state_dict()
for name, param in pretrained_state_dict.items():
    name=name.replace('module.','')
    if name in model_state_dict and model_state_dict[name].shape == param.shape:
        model_state_dict[name] = param
    else:
        print(f"skip loading {name}")
model.load_state_dict(model_state_dict, strict=True)

if use_ddp:
    model = DDP(model, device_ids=[device_id], find_unused_parameters=True)

# step3: Define optimizer and loss
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# For updating learning rate
def update_lr(optimizer, lr):    
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

# if use_distillation:
#     # 蒸馏的损失函数
#     kd_loss = torch.nn.L1Loss()
#     grad_loss = GradInvariantLoss()
# else:
    # midas的损失函数
midas_loss = torch.nn.L1Loss()

weight_L1 = torch.nn.L1Loss(reduction='none')

grad_loss=MultiScaleGradLoss()

# step4: Train the model
total_step = len(train_loader)
curr_lr = learning_rate

start_epoch=-1
cur_iter = 0

# step -1: Initial Tensorboard

for epoch in range(start_epoch + 1 ,num_epochs):
    if use_ddp:
        sampler.set_epoch(epoch)
    total=len(train_loader)
    if use_longtail:
        train_loader.set_epoch(epoch)
    progress_bar = tqdm.tqdm(enumerate(train_loader), total = total)
    all_loss=0
    for i, images in progress_bar:
        
        # images = images
        images, labels, diffs = images

        labels = labels.to(device)
        images = images.to(device)
        diffs = diffs.to(device)

        valmask = (diffs>0)
        thrd_mask = (labels * 255 <= 248) 
        valmask.detach_()
        thrd_mask.detach_()
        cb_mask = valmask & thrd_mask

        # if use_multiscale:
        #     # size = random.choices(input_resolutions, [0.3, 0.45, 0.25])[0]
        #     size = [480,640]
        #     resize_transform = transforms.Resize(size)
        #     images = resize_transform(images)
        #     labels = resize_transform(labels)
        outputs = model(images)
        # print('labels_type = ', type(labels))
        # print('outputs_type = ', type(outputs))

        loss_midas = midas_loss(outputs*255, labels*255)
        loss_grad = grad_loss(outputs*255, labels*255)

        # wieght_loss=weight_L1(outputs*255, labels*255)*diff
        # wieght_loss=wieght_loss.mean()
        # print(loss_grad,loss_midas)

        floss = loss_midas*thrd_mask + loss_grad*thrd_mask
        gloss = midas_loss(outputs[cb_mask]*255, labels[cb_mask]*255)
        '''可以修改成动态参数loss，也可以修改成maskloss'''
        loss = floss + 0.3*gloss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if device_id == 0:
            progress_bar.set_description("Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}"
                .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
        all_loss+=loss.item()
    # Save model
    ave_loss=all_loss/total
    if device_id == 0:
        checkpoint = {"net": model.state_dict(), 'optimizer':optimizer.state_dict(),"epoch": epoch}
        torch.save(checkpoint, save_path+"/"+"ckpt_best_"+str(epoch)+".pth")
        torch.save(model.state_dict(), save_path+"/"+str(epoch)+"-"+str(ave_loss)[0:6]+".pth")

这个代码报错Traceback (most recent call last):
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] File "train_disprefine_maskloss_copy.py", line 208, in
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] loss.backward()
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] File "/opt/conda/envs/zyh/lib/python3.8/site-packages/torch/_tensor.py", line 396, in backward
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] File "/opt/conda/envs/zyh/lib/python3.8/site-packages/torch/autograd/__init__.py", line 166, in backward
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] grad_tensors_ = _make_grads(tensors, grad_tensors_, is_grads_batched=False)
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] File "/opt/conda/envs/zyh/lib/python3.8/site-packages/torch/autograd/__init__.py", line 67, in _make_grads
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] raise RuntimeError("grad can be implicitly created only for scalar outputs")
2024-11-30 01:42:02	2024-11-30 01:42:02.000 [INFO] [Driver] RuntimeError: grad can be implicitly created only for scalar outputs