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'? about u-2-net HOT 4 CLOSED

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

should work

from u-2-net.

u2net_train.py

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

from u-2-net.

thank you, I will try.

from u-2-net.

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?

from u-2-net.