forked from ryoungj/ZO-L2L
-
Notifications
You must be signed in to change notification settings - Fork 1
/
trivial.py
48 lines (37 loc) · 1.7 KB
/
trivial.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import torch
from optimizee import Optimizee
from torch.nn import MSELoss, Parameter
class SimpleConvexModel(Optimizee):
def __init__(self, dim, r=1, cuda=False, gpu_num=None, dtype=torch.float64):
super(SimpleConvexModel, self).__init__()
self.dim = dim
self.r = r
self.cuda = cuda
self.gpu_num = gpu_num
self.dtype = dtype
self.loss_func = MSELoss(reduction='mean')
# Uniform sampling on (-1, 1)
weight = 2 * torch.rand(self.dim, requires_grad=True, dtype=self.dtype) - 1
self.weight = Parameter(weight)
self.v = 2 * torch.rand(self.dim, requires_grad=False, dtype=self.dtype) - 1
if self.cuda:
self.weight = Parameter(self.weight.cuda(self.gpu_num))
self.v = self.v.cuda(self.gpu_num)
def re_initialize(self):
# Generate a new quadratic optimization problem
weight = 2 * torch.rand(self.dim, requires_grad=True, dtype=self.dtype) - 1
self.weight = Parameter(weight)
self.v = 2 * torch.rand(self.dim, requires_grad=False, dtype=self.dtype) - 1
if self.cuda:
self.weight = Parameter(self.weight.cuda(self.gpu_num))
self.v = self.v.cuda(self.gpu_num)
# def reset(self):
# # Uniform sampling on (-1, 1)
# weight = 2 * torch.rand(self.dim, requires_grad=True, dtype=self.dtype) - 1
# self.weight = Parameter(weight)
# self.v = 2 * torch.rand(self.dim, requires_grad=False, dtype=self.dtype) - 1
# if self.cuda:
# self.weight = Parameter(self.weight.cuda(self.gpu_num))
# self.v = self.v.cuda(self.gpu_num)
def loss(self):
return self.loss_func(self.weight, self.v)