# 加载模型
sys.path.append(os.path.join(".", experiment_path, experiment_name))
if arch == "Custom":
from custom import Custom
model = Custom(cfg=cfg.model)
model = torch.nn.DataParallel(model, list(range(
torch.cuda.device_count()))).to(device)
else:
raise NotImplementedError
# 建立tensorboard的实例
from tensorboardX import SummaryWriter
writer = SummaryWriter(os.path.join(".", board_path, experiment_name))
# 建立优化器
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer,
milestones=[10, 20, 30, 40],
gamma=0.5)
# 建立loss
loss_L1 = nn.MSELoss().to(device)
loss_L2 = nn.L1Loss().to(device)
# crossEntropy
loss_BCE = nn.BCELoss().to(device)
# SSIM
loss_SSIM = SSIM(window_size=11, size_average=True)
# L1 + L2 mean
loss_L1_L2 = L1_L2_Loss().to(device)
# 训练的部分
sys.path.append(os.path.join(".", experiment_path, experiment_name))
if arch == "Custom":
from custom import Custom
model = Custom(cfg=cfg.model)
model = model.to(device)
model = torch.nn.DataParallel(model, list(range(
torch.cuda.device_count()))).to(device)
else:
raise NotImplementedError
# 建立tensorboard的实例
from tensorboardX import SummaryWriter
writer = SummaryWriter(os.path.join(".", board_path, experiment_name))
# 建立优化器
optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer,
milestones=[10, 20, 30, 40],
gamma=0.5)
# 建立loss
loss_L1 = nn.MSELoss().to(device)
loss_L2 = nn.L1Loss().to(device)
# crossEntropy
loss_BCE = nn.BCELoss().to(device)
# SSIM
loss_SSIM = SSIM(window_size=11, size_average=True)
loss_L1_L2 = L1_L2_Loss().to(device)
# 训练的部分
