54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import torch
|
|
from torch import optim
|
|
from torch.utils.data import DataLoader
|
|
from data import *
|
|
from net import *
|
|
|
|
from torchvision.utils import save_image
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
weight_path = r'/Users/hanzhangma/Nextcloud/mhz/Study/SS24/MasterThesis/UNet/params/unet.pth'
|
|
data_path = r'/Users/hanzhangma/Document/DataSet/VOC2007'
|
|
save_path = r'/Users/hanzhangma/Nextcloud/mhz/Study/SS24/MasterThesis/Unet/train_image'
|
|
|
|
if __name__ == '__main__':
|
|
data_loader = DataLoader(MyDataset(data_path), batch_size= 4, shuffle=True)
|
|
|
|
net = UNet().to(device)
|
|
if os.path.exists(weight_path):
|
|
net.load_state_dict(torch.load(weight_path))
|
|
print('successful load weight!')
|
|
else:
|
|
print('Failed on load weight!')
|
|
|
|
opt = optim.Adam(net.parameters())
|
|
loss_fun = nn.BCELoss()
|
|
|
|
epoch=1
|
|
|
|
while True:
|
|
for i,(image,segment_image) in enumerate(data_loader):
|
|
image, segment_image = image.to(device), segment_image.to(device)
|
|
|
|
out_image = net(image)
|
|
train_loss = loss_fun(out_image, segment_image)
|
|
|
|
opt.zero_grad()
|
|
train_loss.backward()
|
|
opt.step() # 更新梯度
|
|
|
|
if i%5 ==0 :
|
|
print(f'{epoch} -- {i} -- train loss ===>> {train_loss.item()}')
|
|
|
|
if i % 50 == 0:
|
|
torch.save(net.state_dict(), weight_path)
|
|
|
|
_image = image[0]
|
|
_segment_image = segment_image[0]
|
|
_out_image = out_image[0]
|
|
|
|
img = torch.stack([_image, _segment_image, _out_image], dim=0)
|
|
save_image(img, f'{save_path}/{i}.png')
|
|
|
|
epoch += 1
|