Why the output of my pytorch trained VIT model is very different from the model before being saved?

Viewed 11

I have trained the VIT model (for binary classification) using pytorch distributeddataparallel, the initialized code is below:

net=timm.create_model(model_name=MODEL_NAME,num_classes=NUM_CLASS,pretrained=True,drop_rate=0.1)
loss=nn.CrossEntropyLoss(reduction='none')
net=net.cuda()
net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
optimizer = optim.AdamW(net.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
net,optimizer=amp.initialize(net,optimizer,opt_level="O1")
net = nn.parallel.DistributedDataParallel(net, device_ids=[local_rank],find_unused_parameters=True)

the main training code is below:

loss.sum().backward()
optimizer.step()
scheduler.step(epoch * num_batch + i)
optimizer.zero_grad()

and I save the model every ten epochs like this:

model.eval()
torch.save(model.module.state_dict(),
f'./epoch_model/epoch{epoch}_model_stat_dict.pth')

during the last epoch, the valid acc table is below:

Epoch [300/302] valid_Data: Class 0:  Precision: 0.7915, Recall: 0.9033, Specificity: 0.7450, F1:0.8437
Epoch [300/302] valid_Data: Class 1:  Precision: 0.8779, Recall: 0.7450, Specificity: 0.9033, F1:0.8060

Fusion Matrix:
[243.  26.]
[ 64. 187.]

Epoch [300/302] train loss: 0.0059 acc: 0.9980 kappa: 0.9960valid loss: 1.0297 acc: 0.8269 kappa: 0.6517  [Time: 6.2269]

the output of the model are like this :

[[-0.3752,  0.7028]]

but when i load the model for evaluation with the code like this:

    net =timm.create_model('vit_small_patch16_224',num_classes=2,pretrained=True,drop_rate=0.1)
    net=net.cuda()
    net.load_state_dict(torch.load(
        '../epoch_model/epoch200_model_stat_dict.pth',map_location='cuda:{}'.format(local_rank)))
    net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
    net=amp.initialize(models=net,opt_level="O1")
    net = nn.parallel.DistributedDataParallel(net, device_ids=[local_rank],find_unused_parameters=True)
.........
    model.eval()
    for i,(inputs, labels) in enumerate(dataloader):
        if use_gpu:
            inputs, labels = inputs.cuda(), labels.cuda()
        else:
            inputs, labels = Variable(inputs), Variable(labels)
        # model.eval()
        with torch.no_grad():
            outputs = model(inputs)

the outputs are very different!They are ten times larger than before and the second preditions are all negetive like this :

tensor([[ 8.5312, -9.0859],
        [ 8.4375, -9.0000],
        [ 8.4062, -8.9766],
        [ 8.9062, -9.4375],
        [ 8.5625, -9.1172],
        [ 8.3359, -8.9062],
        [ 8.2109, -8.7891],
        [ 8.3516, -8.9219],
        [ 8.3047, -8.8750],
        [ 8.5234, -9.0781],
        [ 8.2578, -8.8359],
        [ 8.1250, -8.7109],
        [ 8.3828, -8.9531],
        [ 8.6797, -9.2344],
        [ 8.3281, -8.8906],
        [ 8.6953, -9.2344],
        [ 8.7969, -9.3281],
        [ 8.6484, -9.1953],
        [ 8.3438, -8.9141],

Why is this happening?

related information:

system:ubuntu-server 16
pytorch 1.11
cuda 11.4
nvidia-driver 470.141
gpu: rtx 3090 *1
0 Answers
Related