Continue training, train my own model with U-2-Net, in between due to some reasons the training was interrupted or I want to strengthen an existing model, what should I do? Can you provide a 'Continue-training.py'?

[panpengfei21]

Continue training, train my model with U-2-Net, in between due to some reasons the training was interrupted or I want to strengthen an existing model, what should I do? Can you provide a 'Continue-training.py'?

[ghost]

u2net_train.py

if torch.cuda.is_available():
        net.load_state_dict(torch.load(your model path))
        net.cuda()

[panpengfei21]

thank you, I will try.

[panpengfei21]

u2net_train.py

if torch.cuda.is_available():
        net.load_state_dict(torch.load(your model path))
        net.cuda()

My computer is MacBook Pro, so torch.cuda.is_available() is false. I change your code like this:

    checkpoint = torch.load("/xxx/xxx/xxx.pth", map_location='cpu')
    net.load_state_dict(checkpoint)
    if torch.cuda.is_available():
        net.cuda()

Is it OK?

[ghost]

double slash "//xxx//xxx//xxx.pth"

should work