1012 The Best Rank - 模拟

rank有名次的意思,记住了吗!

qwq好吧我第一次都没懂题,就是模拟,附上丑陋代码

还有一个教训是:用ios::sync_with_stdio(false);加速的时候,就不要在用scanf和printf啦!这次就出现了错误

具体解释见这个博客~:https://www.cnblogs.com/cytus/p/7763569.html

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<string>
#include<cstring>
#include<queue>
#include<cmath>
#include<set>
#include<map>
#define ll long long
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
using namespace std;
const int N=2010;
struct A{
    string id;
    int C,M,E,A;
}stu[N];
int c[N],mm[N],e[N],a[N];
map<string,int>ran;
map<string,char>ke;
int main(){
    int n,m;
    ios::sync_with_stdio(false);
    cin>>n>>m;
    string s;
    for(int i=0;i<n;i++){
        cin>>stu[i].id;
        cin>>stu[i].C>>stu[i].M>>stu[i].E;
        c[i]=stu[i].C;mm[i]=stu[i].M;e[i]=stu[i].E;
        stu[i].A=round((double)(stu[i].C+stu[i].M+stu[i].E)/3);
        a[i]=stu[i].A;
    }
    sort(a,a+n);
    sort(c,c+n);
    sort(mm,mm+n);
    sort(e,e+n);
    for(int i=0;i<n;i++){//a,c,m,e
        int r1=upper_bound(a,a+n,stu[i].A)-a;
        r1=n-r1+1;
        ran.insert(pair<string,int>(stu[i].id,r1));
        ke.insert(pair<string,char>(stu[i].id,'A'));
        if(r1==1)continue;
        int r2=upper_bound(c,c+n,stu[i].C)-c;
        r2=n-r2+1;
        if(r2<ran[stu[i].id]){
            ran[stu[i].id]=r2;
            ke[stu[i].id]='C';
        }
        int r3=upper_bound(mm,mm+n,stu[i].M)-mm;
        r3=n-r3+1;
        if(r3<ran[stu[i].id]){
            ran[stu[i].id]=r3;
            ke[stu[i].id]='M';
        }
        int r4=upper_bound(e,e+n,stu[i].E)-e;
        r4=n-r4+1;
        if(r4<ran[stu[i].id]){
            ran[stu[i].id]=r4;
            ke[stu[i].id]='E';
        }
    }
    while(m--){
        cin>>s;
        if(ran.find(s)==ran.end())cout<<"N/A"<<endl;
        else cout<<ran[s]<<" "<<ke[s]<<endl;
    }
}

 

# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # -------------------------------------------------------- # References: # DeiT: https://github.com/facebookresearch/deit # BEiT: https://github.com/microsoft/unilm/tree/master/beit # -------------------------------------------------------- import argparse import datetime import json import numpy as np import os import time from pathlib import Path import torch import torch.backends.cudnn as cudnn import tqdm from torch.utils.tensorboard import SummaryWriter import timm # assert timm.__version__ == "0.3.2" # version check from timm.models.layers import trunc_normal_ from timm.data.mixup import Mixup from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy import util.lr_decay as lrd import util.misc as misc from util.datasets import build_dataset from util.pos_embed import interpolate_pos_embed from util.misc import NativeScalerWithGradNormCount as NativeScaler import models_mae import models_vit from engine_finetune import train_one_epoch import MyDataset import torchvision.transforms as transforms import utils import torch.nn as nn from torch.autograd import Variable import wandb import scipy.io as io from utils_sig import * def get_args_parser(): parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False) parser.add_argument('--batch_size', default=64, type=int, help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus') parser.add_argument('--epochs', default=20, type=int) parser.add_argument('--accum_iter', default=1, type=int, help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)') # Model parameters parser.add_argument('--model', default='vit_base_patch16', type=str, metavar='MODEL', help='Name of model to train') parser.add_argument('--input_size', default=224, type=int, help='images input size') parser.add_argument('--drop_path', type=float, default=0.1, metavar='PCT', help='Drop path rate (default: 0.1)') # Optimizer parameters parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM', help='Clip gradient norm (default: None, no clipping)') parser.add_argument('--weight_decay', type=float, default=0.05, help='weight decay (default: 0.05)') parser.add_argument('--lr', type=float, default=None, metavar='LR', help='learning rate (absolute lr)') parser.add_argument('--blr', type=float, default=1e-3, metavar='LR', help='base learning rate: absolute_lr = base_lr * total_batch_size / 256') parser.add_argument('--layer_decay', type=float, default=0.75, help='layer-wise lr decay from ELECTRA/BEiT') parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR', help='lower lr bound for cyclic schedulers that hit 0') parser.add_argument('--warmup_epochs', type=int, default=50, metavar='N', help='epochs to warmup LR') # Augmentation parameters parser.add_argument('--color_jitter', type=float, default=None, metavar='PCT', help='Color jitter factor (enabled only when not using Auto/RandAug)') parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME', help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)'), parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)') # * Random Erase params parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT', help='Random erase prob (default: 0.25)') parser.add_argument('--remode', type=str, default='pixel', help='Random erase mode (default: "pixel")') parser.add_argument('--recount', type=int, default=1, help='Random erase count (default: 1)') parser.add_argument('--resplit', action='store_true', default=False, help='Do not random erase first (clean) augmentation split') # * Mixup params parser.add_argument('--mixup', type=float, default=0, help='mixup alpha, mixup enabled if > 0.') parser.add_argument('--cutmix', type=float, default=0, help='cutmix alpha, cutmix enabled if > 0.') parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None, help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)') parser.add_argument('--mixup_prob', type=float, default=1.0, help='Probability of performing mixup or cutmix when either/both is enabled') parser.add_argument('--mixup_switch_prob', type=float, default=0.5, help='Probability of switching to cutmix when both mixup and cutmix enabled') parser.add_argument('--mixup_mode', type=str, default='batch', help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"') # * Finetuning params parser.add_argument('--finetune', default='/home/emo/PycharmProjects/rPPG-MAE-main/pretrained_model/STMap_UBFC-100.pth', help='finetune from checkpoint') parser.add_argument('--global_pool', action='store_true', default= False) # parser.set_defaults(global_pool=True) parser.add_argument('--cls_token', action='store_false', dest='global_pool', help='Use class token instead of global pool for classification') # Dataset parameters parser.add_argument('--nb_classes', default=224, type=int, help='number of the classification types') parser.add_argument('--output_dir', default='/home/emo/PycharmProjects/rPPG-MAE-main/finetune_log/my_UBFC_VV', help='path where to save, empty for no saving') parser.add_argument('--log_dir', default='/home/emo/PycharmProjects/rPPG-MAE-main/finetune_log/my_UBFC_VV', help='path where to tensorboard log') parser.add_argument('--device', default='cuda', help='device to use for training / testing') parser.add_argument('--seed', default=0, type=int) parser.add_argument('--resume', default='', help='resume from checkpoint') parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--eval', action='store_true', help='Perform evaluation only') parser.add_argument('--dist_eval', action='store_true', default=False, help='Enabling distributed evaluation (recommended during training for faster monitor') parser.add_argument('--num_workers', default=8, type=int) parser.add_argument('--pin_mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.') parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem') parser.set_defaults(pin_mem=True) # distributed training parameters parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') parser.add_argument('--local_rank', default=-1, type=int) parser.add_argument('--dist_on_itp', action='store_true') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') parser.add_argument('--distributed', action='store_true') parser.add_argument('--dataname', type=str, default="UBFC-PHYS", help='log and save model name') parser.add_argument('--STMap_name1', type=str, default="STMap.png", help='log and save model name') parser.add_argument('--STMap_name2', type=str, default="STMap_YUV_Align_CSI_CHROM.png", help='log and save model name') parser.add_argument('-n', '--frames_num', dest='frames_num', type=int, default=224, help='the num of frames') parser.add_argument('-fn', '--fold_num', type=int, default=5, help='fold_num', dest='fold_num') parser.add_argument('-fi', '--fold_index', type=int, default=0, help='fold_index:0-fold_num', dest='fold_index') parser.add_argument('--log', type=str, default="supervise_VIT_VIPL_LossCrossEntropy", help='log and save model name') parser.add_argument('--loss_type', type=str, default="rppg", help='loss type') parser.add_argument('-rD', '--reData', dest='reData', type=int, default=0, help='re Data') parser.add_argument('--in_chans', type=int, default=3) parser.add_argument('--semi', type=str, default='') return parser def main(args): # misc.init_distributed_mode(args) if args.dataname=='VIPL': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/VIPL_new' # change to your own path. # saveRoot = r'/scratch/project_2006419/data/VIPL_Index/fs_VIPL_STMap' + str(args.fold_num) + str(args.fold_index) # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/VIPL_finetune' # change to your own path. if args.dataname=='PURE': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/Original/PURE' # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/PURE_finetune' # change to your own path. if args.dataname=='UBFC': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/Original/UBFC-rPPG' # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/UBFC_finetune' # change to your own path. if args.dataname == 'UBFC-PHYS': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/Original/UBFC-PHYS' # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/UBFC-PHYS_Train' # change to your own path. if args.dataname == 'VV': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/Original/VV' # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/VV_finetune' # change to your own path. if args.dataname == 'V4V': fileRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/Original/V4V' # change to your own path. saveRoot = r'/home/emo/PycharmProjects/rPPG-MAE-main/Data/V4V_finetune' # change to your own path. # wandb.init(project='rppg-mae'+ args.dataname) # wandb.config = { # "epochs": args.epochs, # "batch_size": args.batch_size # } best_mae = 20 frames_num = args.frames_num dataname = args.dataname fold_num = args.fold_num normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) toTensor = transforms.ToTensor() # resize = transforms.Resize(size=(64, frames_num)) resize = transforms.Resize(size=(frames_num, frames_num)) print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__)))) print("{}".format(args).replace(', ', ',\n')) device = torch.device(args.device) # device = 'cpu' # fix the seed for reproducibility seed = args.seed + misc.get_rank() torch.manual_seed(seed) np.random.seed(seed) cudnn.benchmark = True # 数据集 if args.reData == 1: if args.semi: test_index, train_index, semi_withlabel_index, semi_withoutlabel_index = MyDataset.CrossValidation_semi( fileRoot, fold_num, fold_index=0, semi=2, semi_index=0) semi_with = MyDataset.getIndex(fileRoot, semi_withlabel_index, saveRoot + '_1Train50%', 'STMap.png', 5, frames_num) semi_without = MyDataset.getIndex(fileRoot, semi_withoutlabel_index, saveRoot + '_2Train50%', 'STMap.png', 5, frames_num) else: test_index, train_index = MyDataset.CrossValidation(fileRoot, fold_num=5, fold_index=0) Train_Indexa = MyDataset.getIndex(fileRoot, train_index, saveRoot + '_Train', 'STMap.png', 5, frames_num) Test_Indexa = MyDataset.getIndex(fileRoot, test_index, saveRoot + '_Test', 'STMap.png', 5, frames_num) if args.semi: dataset_train = MyDataset.Data_DG(root_dir=(saveRoot + '_Train' + args.semi), dataName=dataname, STMap1=args.STMap_name1, STMap2=args.STMap_name2, \ in_chans=args.in_chans, frames_num=frames_num, transform=transforms.Compose([resize, toTensor, normalize])) dataset_val = MyDataset.Data_DG(root_dir=(saveRoot + '_Test'), dataName=dataname, STMap1=args.STMap_name1, STMap2=args.STMap_name2, \ in_chans=args.in_chans, frames_num=frames_num, transform=transforms.Compose([resize, toTensor, normalize])) else: dataset_train = MyDataset.Data_DG(root_dir=(saveRoot + '_Train'), dataName=dataname, STMap1=args.STMap_name1, STMap2=args.STMap_name2, \ in_chans=args.in_chans, frames_num=frames_num, transform=transforms.Compose([resize, toTensor, normalize])) dataset_val = MyDataset.Data_DG(root_dir=(saveRoot + '_Test'), dataName=dataname, STMap1=args.STMap_name1, STMap2=args.STMap_name2, \ in_chans=args.in_chans, frames_num=frames_num, transform=transforms.Compose([resize, toTensor, normalize])) print('trainLen:', len(dataset_train), 'testLen:', len(dataset_val)) print('fold_num:', args.fold_num, 'fold_index', args.fold_index) if args.distributed: num_tasks = misc.get_world_size() global_rank = misc.get_rank() sampler_train = torch.utils.data.DistributedSampler( dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True ) print("Sampler_train = %s" % str(sampler_train)) if args.dist_eval: if len(dataset_val) % num_tasks != 0: print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. ' 'This will slightly alter validation results as extra duplicate entries are added to achieve ' 'equal num of samples per-process.') sampler_val = torch.utils.data.DistributedSampler( dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=True) # shuffle=True to reduce monitor bias else: sampler_val = torch.utils.data.SequentialSampler(dataset_val) else: sampler_train = torch.utils.data.RandomSampler(dataset_train) sampler_val = torch.utils.data.SequentialSampler(dataset_val) os.makedirs(args.log_dir, exist_ok=True) log_writer = SummaryWriter(log_dir=args.log_dir) data_loader_train = torch.utils.data.DataLoader( dataset_train, sampler=sampler_train, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=True, ) data_loader_val = torch.utils.data.DataLoader( dataset_val, sampler=sampler_val, batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=args.pin_mem, drop_last=False ) mixup_fn = None mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None if mixup_active: print("Mixup is activated!") mixup_fn = Mixup( mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax, prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode, label_smoothing=args.smoothing, num_classes=args.nb_classes) model = models_vit.__dict__[args.model]( num_classes=args.nb_classes, drop_path_rate=args.drop_path, global_pool=args.global_pool, in_chans=args.in_chans ) if args.finetune: checkpoint = torch.load(args.finetune, map_location='cpu') print("Load pre-trained checkpoint from: %s" % args.finetune) checkpoint_model = checkpoint['model'] state_dict = model.state_dict() for k in ['head.weight', 'head.bias']: if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape: print(f"Removing key {k} from pretrained checkpoint") del checkpoint_model[k] # interpolate position embedding interpolate_pos_embed(model, checkpoint_model) # load pre-trained model msg = model.load_state_dict(checkpoint_model, strict=False) print(msg) if args.global_pool: assert set(msg.missing_keys) == {'head.weight', 'head.bias', 'fc_norm.weight', 'fc_norm.bias'} else: assert set(msg.missing_keys) == {'head.weight', 'head.bias'} # manually initialize fc layer trunc_normal_(model.head.weight, std=2e-5) model.to(device) model_without_ddp = model n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print("Model = %s" % str(model_without_ddp)) print('number of params (M): %.2f' % (n_parameters / 1.e6)) eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size() if args.lr is None: # only base_lr is specified args.lr = args.blr * eff_batch_size / 256 print("base lr: %.2e" % (args.lr * 256 / eff_batch_size)) print("actual lr: %.2e" % args.lr) print("accumulate grad iterations: %d" % args.accum_iter) print("effective batch size: %d" % eff_batch_size) if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) model_without_ddp = model.module # build optimizer with layer-wise lr decay (lrd) param_groups = lrd.param_groups_lrd(model_without_ddp, args.weight_decay, no_weight_decay_list=model_without_ddp.no_weight_decay(), layer_decay=args.layer_decay ) optimizer = torch.optim.AdamW(param_groups, lr=args.lr) loss_scaler = NativeScaler() if mixup_fn is not None: # smoothing is handled with mixup label transform criterion = SoftTargetCrossEntropy() elif args.smoothing > 0.: criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing) else: criterion = torch.nn.CrossEntropyLoss() print("criterion = %s" % str(criterion)) misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler) print(f"Start training for {args.epochs} epochs") start_time = time.time() # for epoch in tqdm.tqdm(range(args.start_epoch, args.epochs)): for epoch in range(args.start_epoch, args.epochs): if args.distributed: data_loader_train.sampler.set_epoch(epoch) train_stats = train_one_epoch( model, criterion, data_loader_train, optimizer, device, epoch, loss_scaler, args.clip_grad, mixup_fn, log_writer=log_writer, args=args ) if args.output_dir: misc.save_model( args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, epoch=epoch) # test # model.eval() # HR_pr_temp = [] # 测试集所有预测心率 # HR_rel_temp = [] # for step, (data1, bvp, _) in enumerate(data_loader_val): # # data = Variable(data).float().to(device=device) # data1 = Variable(data1).float().to(device=device) # # data2 = Variable(data2).float().to(device=device) # bvp = Variable(bvp).float().to(device=device) # # HR_rel = Variable(HR_rel).float().to(device=device) # Wave = bvp.unsqueeze(dim=1) # STMap = data1[:, :, :, 0:frames_num] # Wave = Wave[:, :, 0:frames_num] # b, _, _ = Wave.size() # # outputs = model(data1) # [B,220] # if args.loss_type == 'rppg': # loss_func_rPPG = utils.P_loss3().to(device) # loss_func_SP = utils.SP_loss(device, low_bound=36, high_bound=240, clip_length=args.frames_num).to( # device) # _, hr_pr = loss_func_SP(outputs.unsqueeze(dim=1), HR_rel) # _, hr_rel = loss_func_SP(Wave, HR_rel) # loss = loss_func_rPPG(outputs.unsqueeze(dim=1), Wave) # HR_pr_temp.extend(hr_pr.data.cpu().numpy()) # HR_rel_temp.extend(hr_rel.data.cpu().numpy()) # if args.loss_type == 'SP': # loss_func_SP = utils.SP_loss(device, low_bound=36, high_bound=240, clip_length=args.frames_num).to( # device) # loss, hr_pre = loss_func_SP(outputs.unsqueeze(dim=1), HR_rel) # HR_pr_temp.extend(hr_pre.data.cpu().numpy()) # HR_rel_temp.extend(HR_rel.data.cpu().numpy()) # print('loss_test: ', loss) # ME, STD, MAE, RMSE, MER, P = utils.MyEval(HR_pr_temp, HR_rel_temp) # wandb.log({"MAE": MAE, 'epoch': epoch}) # if best_mae > MAE: # best_mae = MAE # io.savemat(args.log + '/' + 'HR_pr.mat', {'HR_pr': HR_pr_temp}) # 训练结束后保存着所有EPOCHE里效果最好的预测心率 # io.savemat(args.log + '/' + 'HR_rel.mat', {'HR_rel': HR_rel_temp}) # 保存效果最好的真实心率 # print('save best predict HR') log_stats = {**{f'train_{k}': v for k, v in train_stats.items()}, # **{f'test_{k}': v for k, v in test_stats.items()}, 'epoch': epoch, 'n_parameters': n_parameters} if args.output_dir and misc.is_main_process(): if log_writer is not None: log_writer.flush() with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f: f.write(json.dumps(log_stats) + "\n") total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) if __name__ == '__main__': args = get_args_parser() args = args.parse_args() args.decoder_embed_dim = 128 args.decoder_depth = 8 args.norm_pix_loss = False args.reData = 0 # args.output_dir = '/media/emo/WD_5T/UBFC/UBFC-phys/rppg-mae_log' # args.log_dir = '/media/emo/WD_5T/UBFC/UBFC-phys/rppg-mae_log' args.finetune = '/home/emo/PycharmProjects/rPPG-MAE-main/pretrained_model/STMap_UBFC-100.pth' # args.finetune = '/home/emo/PycharmProjects/rPPG-MAE-main/pretrained_model/STMap_PURE-100.pth' # args.finetune = '/home/emo/PycharmProjects/rPPG-MAE-main/pretrained_model/PC-STMap_VIPL-399.pth' # args.in_chans = 6 #torch.Size([768, 6, 16, 16]) from checkpoint PC-STMap_VIPL-399.pth args.epochs = 200 args.dataname = 'VV' if args.output_dir: Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args) 这个是生成checkpoints.pth文件的代码
最新发布
07-29
### 生成 `.pth` 格式模型文件的代码示例 在 PyTorch 中,保存模型通常使用 `torch.save()` 函数,可以保存整个模型对象,也可以仅保存模型的状态字典(state_dict)。推荐的做法是保存状态字典,因为它更轻量,并且在模型结构发生变化时更灵活[^1]。 以下是一个完整的代码示例,展示如何训练一个简单的模型并将其保存为 `.pth` 文件: ```python import torch import torch.nn as nn import torch.optim as optim # 定义一个简单的神经网络模型 class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.linear = nn.Linear(10, 1) def forward(self, x): return self.linear(x) # 创建模型实例 model = SimpleModel() # 定义损失函数和优化器 criterion = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=0.01) # 模拟训练数据 inputs = torch.randn(100, 10) targets = torch.randn(100, 1) # 训练循环 for epoch in range(10): # 前向传播 outputs = model(inputs) loss = criterion(outputs, targets) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() print(f'Epoch [{epoch+1}/10], Loss: {loss.item():.4f}') # 保存模型的状态字典到 .pth 文件 torch.save(model.state_dict(), 'checkpoints.pth') print("模型已保存为 checkpoints.pth") ``` ### 加载 `.pth` 文件并查看模型内容 可以使用 `torch.load()` 函数加载 `.pth` 文件,并通过 `state_dict()` 恢复模型参数: ```python # 加载模型参数 model = SimpleModel() model.load_state_dict(torch.load('checkpoints.pth')) model.eval() # 设置模型为评估模式 # 查看模型参数 print("模型状态字典:") for param_tensor in model.state_dict(): print(param_tensor, "\t", model.state_dict()[param_tensor].size()) ``` ### 保存完整模型与状态字典的区别 - **保存完整模型**:`torch.save(model, 'complete_model.pth')`,该方式会保存整个模型对象,包括模型结构和参数。适用于快速加载模型并直接使用的情况。 - **保存状态字典**:`torch.save(model.state_dict(), 'state_dict_model.pth')`,该方式仅保存模型的参数,不包含模型结构。加载时需要先定义模型结构,再加载参数。更推荐使用此方式,因为更灵活且占用空间更小[^2]。 ### 校验 `.pth` 文件内容 可以将 `.pth` 文件的内容写入文本文件进行查看,以确认是否保存正确: ```python import torch # 加载模型文件 pthfile = './checkpoints.pth' model = torch.load(pthfile) # 写入文本文件查看内容 with open('weight.txt', 'w') as f: f.write('type:\n') f.write(str(type(model)) + '\n') f.write('length:\n') f.write(str(len(model)) + '\n\n') f.write('key:\n') for k in model.keys(): f.write(k + '\n') f.write('\n') f.write('value:\n') for k in model: f.write(k + ': ' + repr(model[k]) + '\n') ``` ### 总结 生成 `.pth` 文件的核心是使用 `torch.save()` 保存模型的状态字典或完整模型对象。保存状态字典是推荐的做法,因为它在模型结构发生变化时更具灵活性。加载时需要先定义模型结构,再使用 `load_state_dict()` 加载参数[^1]。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值