【pytorch】训练模型第一步，加载模型和指定device

1. torch.device 选择GPU或者CPU

torch.cuda.in_available() 判断cuda是否可用
torch.device(‘cuda’) 将数据转到GPU
torch.device('cpu‘） 将数据转到CPU
torch.device('cuda:0') 指定设备号， 不指定则默认current_device
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

2. timm.create_model 加载模型

model = timm.create_model('resnet18', pretrained=True, 
pretrained_cfg_overlay=dict(file='/home/xiaoxin/Documents/hc/-/bin/pytorch_model.bin'))

3. torch.load() 加载预训练模型

如果使用cpu,则需要添加参数map_location= ‘cpu’，默认使用GPU

checkpoint_path = "0515.pth"
if os.path.exists(checkpoint_path):
    print('load checkpoint:%s'%checkpoint_path)
    checkpoint = torch.load(checkpoint_path, map_location= 'cpu')
    model_ft.load_state_dict(checkpoint, strict=False)

4. 确认模型训练的设备

#确认模型用CDEVICE训练，确认设备的情况下，也可以写成model.cuda()或者model.cpu

model.to(DEVICE)
print(model.get_classifier())