参数更新和三个单条件比较 · 实现代码

%%writefile -a lab11_segmentation.py
def train_model(model,x,target,epochs,lr):
    opt=torch.optim.Adam(model.parameters(),lr=lr);history=[]
    for epoch in range(epochs):
        model.train();losses=[]
        for ids in torch.randperm(len(x)).split(32):
            images=x[ids];labels=target[ids]
            # One geometric decision applies to both image and mask.
            if torch.rand(())<.5:images=images.flip(-1);labels=labels.flip(-1)
            logits=model(images)
            assert logits.shape==(len(ids),3,48,48) and labels.dtype==torch.long
            loss=F.cross_entropy(logits,labels,ignore_index=255)
            opt.zero_grad();loss.backward();opt.step();losses.append(float(loss.detach()))
        history.append(float(np.mean(losses)))
        if epoch==0 or (epoch+1)%5==0:print(type(model).__name__,epoch+1,history[-1],flush=True)
    return history

@torch.no_grad()
def predict(model,x,mode='bilinear'):
    model.eval();return torch.cat([model(b,mode=mode).argmax(1) for b in x.split(32)])

def plot_predictions(images,truth,pred,path,title,indices=None):
    if indices is None:indices=[0,1,2,3]
    palette=np.array([[.04,.04,.05],[.95,.22,.15],[.10,.45,1],[.55,.55,.55]])
    fig,axes=plt.subplots(len(indices),4,figsize=(10,2.6*len(indices)),squeeze=False)
    for row,i in enumerate(indices):
        mask=truth[i].numpy().copy();mask[mask==255]=3
        error=(pred[i]!=truth[i])&(truth[i]!=255)
        data=[images[i].permute(1,2,0).numpy(),palette[mask],palette[pred[i].numpy()],error.numpy()]
        for ax,d,name in zip(axes[row],data,['Input RGB','Target / gray=ignore','Model prediction','Error on valid pixels']):
            ax.imshow(d,cmap='magma',interpolation='nearest');ax.axis('off');ax.set_title(f'{name} #{i}',fontsize=9)
    fig.suptitle(title);fig.tight_layout();fig.savefig(path,dpi=140);plt.close(fig)