mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-10-28 20:30:42 +00:00
add example of self-supervised SimCLR training - V2 (#50)
* add example of self-supervised SimCLR training * simclr v2, replace nvidia dali dataloader * updated * sync to latest code writing style * sync to latest code writing style and modify README * detail README & standardize dataset path
This commit is contained in:
19
examples/simclr_cifar10_data_parallel/models/linear_eval.py
Normal file
19
examples/simclr_cifar10_data_parallel/models/linear_eval.py
Normal file
@@ -0,0 +1,19 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from .Backbone import backbone
|
||||
|
||||
class Linear_eval(nn.Module):
|
||||
|
||||
def __init__(self, model='resnet18', class_num=10, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
self.backbone = backbone(model, **kwargs)
|
||||
self.backbone.requires_grad_(False)
|
||||
self.fc = nn.Linear(self.backbone.output_dim, class_num)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.backbone(x)
|
||||
out = self.fc(out)
|
||||
return out
|
||||
Reference in New Issue
Block a user